Fix for unit test
Code cleanup