Fixed up unit test

Hopefully this will fix failing test.