[UnitTest] Fixed unit test (hopefully).