From d5fc105d4678376ae259dd3a3ecb45fabc4afba0 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 6 May 2025 16:00:08 -0400
Subject: [PATCH 1/2] fix(load-task): add some environment variables normally
 set by the worker

---
 src/taskgraph/docker.py | 11 ++++++++++-
 test/test_docker.py     |  3 +++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index 0bd7d8035..72dbe18ce 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -21,6 +21,7 @@
 from taskgraph.util import docker
 from taskgraph.util.taskcluster import (
     get_artifact_url,
+    get_root_url,
     get_session,
     get_task_definition,
 )
@@ -295,7 +296,15 @@ def load_task(task_id, remove=True):
     image_task_id = task_def["payload"]["image"]["taskId"]
     image_tag = load_image_by_task_id(image_task_id)
 
-    env = task_def["payload"].get("env")
+    # Set some env vars the worker would normally set.
+    env = {
+        "RUN_ID": "0",
+        "TASK_GROUP_ID": task_def.get("taskGroupId", ""),
+        "TASK_ID": task_id,
+        "TASKCLUSTER_ROOT_URL": get_root_url(False),
+    }
+    # Add the task's environment variables.
+    env.update(task_def["payload"].get("env", {}))
 
     envfile = None
     initfile = None
diff --git a/test/test_docker.py b/test/test_docker.py
index db010454d..40bf11611 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -151,6 +151,7 @@ def test_load_task(run_load_task):
         "run",
         "-v",
         re.compile(f"{tempfile.gettempdir()}/tmp.*:/builds/worker/.bashrc"),
+        re.compile(f"--env-file={tempfile.gettempdir()}/tmp.*"),
         "-it",
         "image/tag",
         "bash",
@@ -163,6 +164,8 @@ def test_load_task(run_load_task):
     mocks["subprocess_run"].assert_called_once()
     actual = mocks["subprocess_run"].call_args[0][0]
 
+    print(expected)
+    print(actual)
     assert len(expected) == len(actual)
     for i, exp in enumerate(expected):
         if isinstance(exp, re.Pattern):

From 9daedf78b046cc2fe534069d56bfeccf06063feb Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 6 May 2025 16:00:08 -0400
Subject: [PATCH 2/2] fix(load-task): ensure final bash is invoked as 'worker'
 user

---
 src/taskgraph/docker.py | 2 +-
 test/test_docker.py     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index 72dbe18ce..8dc76b8e9 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -316,7 +316,7 @@ def load_task(task_id, remove=True):
             image_tag,
             "bash",
             "-c",
-            f"{shlex.join(command)} && cd $TASK_WORKDIR && bash",
+            f"{shlex.join(command)} && cd $TASK_WORKDIR && su -p worker",
         ]
 
         if remove:
diff --git a/test/test_docker.py b/test/test_docker.py
index 40bf11611..3bc3d3638 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -158,7 +158,7 @@ def test_load_task(run_load_task):
         "-c",
         "/usr/bin/run-task --repo-checkout=/builds/worker/vcs/repo "
         "--task-cwd=/builds/worker/vcs/repo -- echo 'Task setup complete!\n"
-        "Run `exec-task` to execute the task'\"'\"'s command.' && cd $TASK_WORKDIR && bash",
+        "Run `exec-task` to execute the task'\"'\"'s command.' && cd $TASK_WORKDIR && su -p worker",
     ]
 
     mocks["subprocess_run"].assert_called_once()