Unit testing fixes