fix(load-task): add some environment variables normally set by the worker

Author: ahal

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

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):