Unit test tweaks