Updated unit test framework