Session tests and a little bit of refactoring