Updated unit tests and modified base classes