Partial fix for unit test failure