Fixing an import in the test suite