Unit test corrections