diff --git a/tests/test_domain.py b/tests/test_domain.py
index 326d3dd..a67ff1c 100644
--- a/tests/test_domain.py
+++ b/tests/test_domain.py
@@ -1,4 +1,6 @@
-from simple_shapes.domain import (
+from utils import PROJECT_DIR
+
+from simple_shapes_dataset.domain import (
     Attribute,
     RawText,
     SimpleShapesAttributes,
@@ -7,7 +9,6 @@
     SimpleShapesText,
     Text,
 )
-from utils import PROJECT_DIR
 
 
 def test_image_domain():