Addressing unit test failures

Addressing unit test failures