Refactoring and improvements to unit tests (with some breakage)