More fixes for unit tests

Necessary to get tests to compile again after the latest refactor.