Unit test improvements