Fix tests

The tests were broken by the refactoring. This fixes them.