fix: updated to review and added tests

Author: abhishekmadan30

src/taskgraph/config.py

@@ -28,8 +28,9 @@
         Required("trust-domain"): str,
         Optional(
             "docker-image-kind",
-            description="What kind of docker image to use for the task.",
-        ): [str],
+            default="docker-image",
+            description="Name of the docker image kind (default: docker-image)",
+        ): str,
         Required("task-priority"): optionally_keyed_by(
             "project",
             "level",

src/taskgraph/util/docker.py

@@ -206,33 +206,31 @@ def stream_context_tar(topsrcdir, context_dir, out_file, args=None):
 
 
 @functools.lru_cache(maxsize=None)
-def image_paths(graph_config=None):
+def image_paths(graph_config):
     """Return a map of image name to paths containing their Dockerfile."""
-    if graph_config:
-        config = load_yaml(
-            graph_config.docker_dir,
-            "kinds",
-            graph_config["docker_image_kind"],
-            "kind.yml",
-        )
-    else:
-        config = load_yaml("taskcluster", "kinds", "docker-image", "kind.yml")
+
+    config = load_yaml(
+        graph_config.root_dir,
+        "kinds",
+        graph_config["docker-image-kind"],
+        "kind.yml",
+    )
 
     return {
-        k: os.path.join(IMAGE_DIR, v.get("definition", k))
+        k: os.path.join(graph_config.docker_dir, v.get("definition", k))
         for k, v in config["tasks"].items()
     }
 
 
-def image_path(name, graph_config=None):
+def image_path(name, graph_config):
     paths = image_paths(graph_config)
     if name in paths:
         return paths[name]
-    return os.path.join(IMAGE_DIR, name)
+    return os.path.join(graph_config.docker_dir, name)
 
 
 @functools.lru_cache(maxsize=None)
-def parse_volumes(image, graph_config=None):
+def parse_volumes(image, graph_config):
     """Parse VOLUME entries from a Dockerfile for an image."""
     volumes = set()
 

taskcluster/config.yml

@@ -14,6 +14,8 @@ index:
 
 task-priority: low
 
+docker-image-kind: docker-image
+
 taskgraph:
     register: self_taskgraph:register
     decision-parameters: 'self_taskgraph.custom_parameters:decision_parameters'

test/test_util_docker.py

@@ -14,6 +14,7 @@
 
 import taskcluster_urls as liburls
 
+from taskgraph.config import GraphConfig
 from taskgraph.util import docker
 
 from .mockedopen import MockedOpen
@@ -269,3 +270,81 @@ def test_stream_context_tar(self):
             )
         finally:
             shutil.rmtree(tmp)
+
+    def test_image_paths_with_custom_kind(self):
+        """Test image_paths function with graph_config parameter."""
+        temp_dir = tempfile.mkdtemp()
+        try:
+            # Create the kinds directory structure
+            kinds_dir = os.path.join(temp_dir, "kinds", "docker-test-image")
+            os.makedirs(kinds_dir)
+
+            # Create the kind.yml file with task definitions
+            kind_yml_path = os.path.join(kinds_dir, "kind.yml")
+            with open(kind_yml_path, "w") as f:
+                f.write("tasks:\n")
+                f.write("  test-image:\n")
+                f.write("    definition: test-image\n")
+                f.write("  another-image:\n")
+                f.write("    definition: custom-path\n")
+
+            # Create graph config pointing to our test directory
+            temp_graph_config = GraphConfig(
+                {
+                    "trust-domain": "test-domain",
+                    "docker-image-kind": "docker-test-image",
+                },
+                temp_dir,
+            )
+
+            paths = docker.image_paths(temp_graph_config)
+
+            expected_docker_dir = os.path.join(temp_graph_config.root_dir, "docker")
+            self.assertEqual(
+                paths["test-image"], os.path.join(expected_docker_dir, "test-image")
+            )
+            self.assertEqual(
+                paths["another-image"], os.path.join(expected_docker_dir, "custom-path")
+            )
+        finally:
+            shutil.rmtree(temp_dir)
+
+    def test_parse_volumes_with_graph_config(self):
+        """Test parse_volumes function with graph_config parameter."""
+        temp_dir = tempfile.mkdtemp()
+        try:
+            kinds_dir = os.path.join(temp_dir, "kinds", "docker-test-image")
+            os.makedirs(kinds_dir)
+
+            kind_yml_path = os.path.join(kinds_dir, "kind.yml")
+            with open(kind_yml_path, "w") as f:
+                f.write("tasks:\n")
+                f.write("  test-image:\n")
+                f.write("    definition: test-image\n")
+
+            docker_dir = os.path.join(temp_dir, "docker")
+            os.makedirs(docker_dir)
+
+            image_dir = os.path.join(docker_dir, "test-image")
+            os.makedirs(image_dir)
+
+            dockerfile_path = os.path.join(image_dir, "Dockerfile")
+            with open(dockerfile_path, "wb") as fh:
+                fh.write(b"VOLUME /foo/bar \n")
+                fh.write(b"VOLUME /hello  /world \n")
+
+            test_graph_config = GraphConfig(
+                {
+                    "trust-domain": "test-domain",
+                    "docker-image-kind": "docker-test-image",
+                },
+                temp_dir,
+            )
+
+            volumes = docker.parse_volumes("test-image", test_graph_config)
+
+            expected_volumes = {"/foo/bar", "/hello", "/world"}
+            self.assertEqual(volumes, expected_volumes)
+
+        finally:
+            shutil.rmtree(temp_dir)