Unit test failure fixes