Fix exception handling regression in load_task

src/taskgraph/docker.py

@@ -16,16 +16,15 @@
 from textwrap import dedent
 from typing import Any, Dict, Generator, List, Optional, Union
 
-from requests import HTTPError
-
-from taskgraph.generator import load_tasks_for_kind
+from taskcluster.exceptions import TaskclusterRestFailure
 
 try:
     import zstandard as zstd
 except ImportError as e:
     zstd = e
 
 from taskgraph.config import GraphConfig
+from taskgraph.generator import load_tasks_for_kind
 from taskgraph.transforms import docker_image
 from taskgraph.util import docker, json
 from taskgraph.util.taskcluster import (
@@ -203,8 +202,8 @@ def build_image(
         if parent_id := task_def["payload"].get("env", {}).get("PARENT_TASK_ID"):
             try:
                 status_task(parent_id)
-            except HTTPError as e:
-                if e.response.status_code != 404:
+            except TaskclusterRestFailure as e:
+                if e.status_code != 404:
                     raise
 
                 # Parent id doesn't exist, needs to be re-built as well.

test/data/taskcluster/kinds/docker-image/kind.yml

@@ -1,38 +1,11 @@
-# This Source Code Form is subject to the terms of the Mozilla Public
-# License, v. 2.0. If a copy of the MPL was not distributed with this
-# file, You can obtain one at http://mozilla.org/MPL/2.0/.
 ---
-meta:
-    - &uv_version 0.6.1
-
 loader: taskgraph.loader.transform:loader
 
 transforms:
     - taskgraph.transforms.docker_image:transforms
     - taskgraph.transforms.cached_tasks:transforms
     - taskgraph.transforms.task:transforms
 
-# make a task for each docker-image we might want.  For the moment, since we
-# write artifacts for each, these are whitelisted, but ideally that will change
-# (to use subdirectory clones of the proper directory), at which point we can
-# generate tasks for every docker image in the directory, secure in the
-# knowledge that unnecessary images will be omitted from the target task graph
 tasks:
-    decision:
-        symbol: I(d)
-        parent: run-task
-    fetch:
-        symbol: I(fetch)
-    index-task:
-        symbol: I(idx)
-    python:
-        symbol: I(py)
-        args:
-            PYTHON_VERSIONS: "3.13 3.12 3.11 3.10 3.9 3.8"
-            UV_VERSION: *uv_version
-    run-task:
-        symbol: I(rt)
-        args:
-            UV_VERSION: *uv_version
-    skopeo:
-        symbol: I(skopeo)
+    hello-world: {}
+    hello-world-tag: {}

test/test_docker.py

@@ -1,4 +1,5 @@
 import logging
+import os
 import re
 import tempfile
 
@@ -9,6 +10,17 @@
 from taskgraph.transforms.docker_image import IMAGE_BUILDER_IMAGE
 
 
+@pytest.fixture
+def root_url():
+    return "https://tc.example.com"
+
+
+@pytest.fixture(autouse=True)
+def mock_environ(monkeypatch, root_url):
+    # Ensure user specified environment variables don't interfere with URLs.
+    monkeypatch.setattr(os, "environ", {"TASKCLUSTER_ROOT_URL": root_url})
+
+
 @pytest.fixture(autouse=True, scope="module")
 def mock_docker_path(module_mocker, datadir):
     module_mocker.patch(
@@ -284,6 +296,36 @@ def test_load_task_with_different_image_types(
     mocks["load_image_by_task_id"].assert_called_once_with(image_task_id)
 
 
+def test_load_task_with_local_image(
+    mocker,
+    run_load_task,
+):
+    task_id = "abc"
+    image_task_id = "xyz"
+    task = {
+        "metadata": {"name": "test-task-image-types"},
+        "payload": {
+            "command": [
+                "/usr/bin/run-task",
+                "--task-cwd=/builds/worker",
+                "--",
+                "echo",
+                "test",
+            ],
+            "image": "hello-world",
+        },
+    }
+
+    mocker.patch.object(docker, "find_task_id", return_value=image_task_id)
+
+    ret, mocks = run_load_task(task)
+    assert ret == 0
+
+    mocks["get_task_definition"].assert_called_once_with(task_id)
+    mocks["build_image"].assert_called_once()
+    assert mocks["build_image"].call_args[0][1] == "hello-world"
+
+
 def test_load_task_with_unsupported_image_type(caplog, run_load_task):
     caplog.set_level(logging.DEBUG)
     task = {
@@ -439,7 +481,7 @@ def test_load_task_with_custom_image_registry(mocker, run_load_task, task):
 
 @pytest.fixture
 def run_build_image(mocker):
-    def inner(image_name, save_image=None, context_file=None):
+    def inner(image_name, save_image=None, context_file=None, image_task=None):
         graph_config = GraphConfig(
             {
                 "trust-domain": "test-domain",
@@ -494,17 +536,23 @@ def mock_path_constructor(path_arg):
             "subprocess": mocker.patch.object(docker.subprocess, "run"),
             "shutil_copy": mocker.patch.object(docker.shutil, "copy"),
             "shutil_move": mocker.patch.object(docker.shutil, "move"),
-            "status_task": mocker.patch.object(docker, "status_task"),
             "isdir": mocker.patch.object(docker.os.path, "isdir", return_value=True),
             "getuid": mocker.patch.object(docker.os, "getuid", return_value=1000),
             "getgid": mocker.patch.object(docker.os, "getgid", return_value=1000),
         }
 
         # Mock image task
-        mocks["task"] = mocker.MagicMock()
-        mocks["task"].task = {"payload": {"env": {}}}
+        if not image_task:
+            image_task = mocker.MagicMock()
+            image_task.task = {"payload": {"env": {}}}
+
+        parent_image = mocker.MagicMock()
+        parent_image.task = {"payload": {"env": {}}}
+
+        mocks["image_task"] = image_task
         mocks["load_tasks_for_kind"].return_value = {
-            f"docker-image-{image_name}": mocks["task"]
+            f"docker-image-{image_name}": mocks["image_task"],
+            "docker-image-parent": parent_image,
         }
 
         # Mock subprocess result for docker load
@@ -545,7 +593,42 @@ def test_build_image(run_build_image):
     mocks["load_task"].assert_called_once()
     call_args = mocks["load_task"].call_args
     assert call_args[0][0] == mocks["graph_config"]
-    assert call_args[0][1] == mocks["task"].task
+    assert call_args[0][1] == mocks["image_task"].task
+    assert call_args[1]["custom_image"] == IMAGE_BUILDER_IMAGE
+    assert call_args[1]["interactive"] is False
+    assert "volumes" in call_args[1]
+
+    # Verify docker load was called
+    mocks["subprocess"].assert_called_once()
+    docker_load_args = mocks["subprocess"].call_args[0][0]
+    assert docker_load_args[:3] == ["docker", "load", "-i"]
+
+    assert result == "hello-world:latest"
+
+
+def test_build_image_with_parent(mocker, responses, root_url, run_build_image):
+    parent_task_id = "abc"
+    responses.get(f"{root_url}/api/queue/v1/task/{parent_task_id}/status")
+
+    # Test building image that has a parent image
+    image_task = mocker.MagicMock()
+    image_task.task = {"payload": {"env": {"PARENT_TASK_ID": parent_task_id}}}
+    result, mocks = run_build_image("hello-world", image_task=image_task)
+    assert result == "hello-world:latest"
+
+    # Verify the graph generation call
+    mocks["load_tasks_for_kind"].assert_called_once_with(
+        {"do_not_optimize": ["docker-image-hello-world"]},
+        "docker-image",
+        graph_attr="morphed_task_graph",
+        write_artifacts=True,
+    )
+
+    # Verify load-task called (to invoke image_builder)
+    mocks["load_task"].assert_called_once()
+    call_args = mocks["load_task"].call_args
+    assert call_args[0][0] == mocks["graph_config"]
+    assert call_args[0][1] == mocks["image_task"].task
     assert call_args[1]["custom_image"] == IMAGE_BUILDER_IMAGE
     assert call_args[1]["interactive"] is False
     assert "volumes" in call_args[1]
@@ -555,8 +638,44 @@ def test_build_image(run_build_image):
     docker_load_args = mocks["subprocess"].call_args[0][0]
     assert docker_load_args[:3] == ["docker", "load", "-i"]
 
+
+def test_build_image_with_parent_not_found(
+    mocker, responses, root_url, run_build_image
+):
+    parent_task_id = "abc"
+    responses.get(f"{root_url}/api/queue/v1/task/{parent_task_id}/status", status=404)
+
+    # Test building image that uses DOCKER_IMAGE_PARENT
+    image_task = mocker.MagicMock()
+    image_task.task = {"payload": {"env": {"PARENT_TASK_ID": parent_task_id}}}
+    image_task.dependencies = {"parent": "docker-image-parent"}
+    result, mocks = run_build_image("hello-world", image_task=image_task)
     assert result == "hello-world:latest"
 
+    # Verify the graph generation call
+    assert mocks["load_tasks_for_kind"].call_count == 2
+    assert mocks["load_tasks_for_kind"].call_args_list[0] == (
+        ({"do_not_optimize": ["docker-image-hello-world"]}, "docker-image"),
+        {"graph_attr": "morphed_task_graph", "write_artifacts": True},
+    )
+    assert mocks["load_tasks_for_kind"].call_args_list[1] == (
+        ({"do_not_optimize": ["docker-image-parent"]}, "docker-image"),
+        {"graph_attr": "morphed_task_graph", "write_artifacts": True},
+    )
+
+    # Verify load-task called (to invoke image_builder)
+    assert mocks["load_task"].call_count == 2
+    call_args = mocks["load_task"].call_args_list[0]
+    assert call_args[0][0] == mocks["graph_config"]
+    assert call_args[1]["custom_image"] == IMAGE_BUILDER_IMAGE
+    assert call_args[1]["interactive"] is False
+    assert "volumes" in call_args[1]
+
+    # Verify docker load was called
+    mocks["subprocess"].assert_called_once()
+    docker_load_args = mocks["subprocess"].call_args[0][0]
+    assert docker_load_args[:3] == ["docker", "load", "-i"]
+
 
 def test_build_image_with_save_image(run_build_image):
     save_path = "/path/to/save.tar"