Even more unit test and code cleanup