diff --git a/tests/test_query_relations.py b/tests/test_query_relations.py
index 32fa894ec1..2521ef01de 100644
--- a/tests/test_query_relations.py
+++ b/tests/test_query_relations.py
@@ -31,8 +31,12 @@
 import re
 import unittest
 
-from lsst.daf.butler import DataCoordinateQueryResults, DatasetRefQueryResults, DimensionRecordQueryResults
 from lsst.daf.butler.registry import MissingSpatialOverlapError, RegistryConfig, _RegistryFactory
+from lsst.daf.butler.registry.queries import (
+    DataCoordinateQueryResults,
+    DatasetQueryResults,
+    DimensionRecordQueryResults,
+)
 from lsst.daf.butler.transfers import YamlRepoImportBackend
 
 TESTDIR = os.path.abspath(os.path.dirname(__file__))
@@ -99,7 +103,7 @@ def setUpClass(cls) -> None:
     def assert_relation_str(
         self,
         expected: str,
-        *results: DataCoordinateQueryResults | DimensionRecordQueryResults | DatasetRefQueryResults,
+        *results: DataCoordinateQueryResults | DimensionRecordQueryResults | DatasetQueryResults,
     ) -> None:
         """Assert that checks that one or more registry
         queries have relation trees that match the given string.