More unit test work