Update unit tests to work with merged codebase