Corrected unit tests for new code