diff --git a/tests/hierarchy/get_test_dirs.py b/tests/hierarchy/get_test_dirs.py
new file mode 100644
index 00000000..e9165ec8
--- /dev/null
+++ b/tests/hierarchy/get_test_dirs.py
@@ -0,0 +1,21 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+from pathlib import Path
+
+from tests.hierarchy.images_hierarchy import IMAGE_PARENT
+
+THIS_DIR = Path(__file__).parent.resolve()
+IMAGE_SPECIFIC_TESTS_DIR = THIS_DIR.parent / "image_specific_tests"
+
+assert IMAGE_SPECIFIC_TESTS_DIR.exists(), f"{IMAGE_SPECIFIC_TESTS_DIR} does not exist."
+
+
+def get_test_dirs(image: str | None) -> list[Path]:
+    if image is None:
+        return []
+
+    test_dirs = get_test_dirs(IMAGE_PARENT[image])
+    current_test_dir = IMAGE_SPECIFIC_TESTS_DIR / image
+    assert current_test_dir.exists(), f"{current_test_dir} does not exist."
+    test_dirs.append(current_test_dir)
+    return test_dirs
diff --git a/tests/hierarchy/images_hierarchy.py b/tests/hierarchy/images_hierarchy.py
index 433bb821..180fe4ea 100644
--- a/tests/hierarchy/images_hierarchy.py
+++ b/tests/hierarchy/images_hierarchy.py
@@ -1,15 +1,9 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
-from pathlib import Path
-
-THIS_DIR = Path(__file__).parent.resolve()
-IMAGE_SPECIFIC_TESTS_DIR = THIS_DIR.parent / "image_specific_tests"
-
-assert IMAGE_SPECIFIC_TESTS_DIR.exists(), f"{IMAGE_SPECIFIC_TESTS_DIR} does not exist."
 
 # Please, take a look at the hierarchy of the images here:
 # https://jupyter-docker-stacks.readthedocs.io/en/latest/using/selecting.html#image-relationships
-_IMAGE_PARENT = {
+IMAGE_PARENT = {
     "docker-stacks-foundation": None,
     "base-notebook": "docker-stacks-foundation",
     "minimal-notebook": "base-notebook",
@@ -22,14 +16,3 @@ _IMAGE_PARENT = {
     "pyspark-notebook": "scipy-notebook",
     "all-spark-notebook": "pyspark-notebook",
 }
-
-
-def get_test_dirs(image: str | None) -> list[Path]:
-    if image is None:
-        return []
-
-    test_dirs = get_test_dirs(_IMAGE_PARENT[image])
-    current_test_dir = IMAGE_SPECIFIC_TESTS_DIR / image
-    assert current_test_dir.exists(), f"{current_test_dir} does not exist."
-    test_dirs.append(current_test_dir)
-    return test_dirs
diff --git a/tests/image_specific_tests/docker-stacks-foundation/test_units.py b/tests/image_specific_tests/docker-stacks-foundation/test_units.py
index 9e34e92e..04d4225f 100644
--- a/tests/image_specific_tests/docker-stacks-foundation/test_units.py
+++ b/tests/image_specific_tests/docker-stacks-foundation/test_units.py
@@ -2,7 +2,7 @@
 # Distributed under the terms of the Modified BSD License.
 import logging
 
-from tests.hierarchy.images_hierarchy import get_test_dirs
+from tests.hierarchy.get_test_dirs import get_test_dirs
 from tests.utils.tracked_container import TrackedContainer
 
 LOGGER = logging.getLogger(__name__)
diff --git a/tests/run_tests.py b/tests/run_tests.py
index 4fad73c8..2de04619 100755
--- a/tests/run_tests.py
+++ b/tests/run_tests.py
@@ -6,7 +6,7 @@ import logging
 
 import plumbum
 
-from tests.hierarchy.images_hierarchy import get_test_dirs
+from tests.hierarchy.get_test_dirs import get_test_dirs
 
 python3 = plumbum.local["python3"]