Unit test fixing 2