cleaned up unit test
Improved test coverage