ci: implement an optimization for the integration tests

taskcluster/fxci_config_taskgraph/__init__.py

@@ -7,6 +7,7 @@ def register(graph_config):
     # Import sibling modules, triggering decorators in the process
     _import_modules(
         [
+            "optimizations",
             "target_tasks",
         ]
     )

taskcluster/fxci_config_taskgraph/optimizations.py

@@ -0,0 +1,45 @@
+import os
+import subprocess
+from functools import cache
+
+from taskgraph.optimize.base import OptimizationStrategy, register_strategy
+from taskgraph.optimize.strategies import SkipUnlessChanged
+
+from fxci_config_taskgraph.util.constants import FIREFOXCI_ROOT_URL
+
+
+@register_strategy("integration-test")
+class IntegrationTestStrategy(OptimizationStrategy):
+    @cache
+    def _get_modified_worker_pools(self) -> set[str]:
+        cmd = [
+            "tc-admin",
+            "diff",
+            "--environment",
+            "firefoxci",
+            "--resources",
+            "worker_pools",
+            "--ids-only",
+        ]
+        env = os.environ.copy()
+        env["TASKCLUSTER_ROOT_URL"] = FIREFOXCI_ROOT_URL
+        proc = subprocess.run(cmd, capture_output=True, text=True, env=env)
+        lines = [line for line in proc.stdout.splitlines() if line.startswith("!")]
+
+        worker_pools = set()
+        for line in lines:
+            line = line.split()[1]
+            assert line.startswith("WorkerPool=")
+            line = line.split("=", 1)[1]
+            worker_pools.add(line)
+
+        return worker_pools
+
+    def should_remove_task(self, task, params, files_changed):
+        task_queue_id = f"{task.task['provisionerId']}/{task.task['workerType']}"
+        if task_queue_id in self._get_modified_worker_pools():
+            return False
+
+        # If the worker type wasn't impacted, defer to `skip-unless-changed`.
+        opt = SkipUnlessChanged()
+        return opt.should_remove_task(task, params, files_changed)

taskcluster/fxci_config_taskgraph/transforms/integration_test.py

@@ -182,6 +182,13 @@ def make_integration_test_description(task_def: dict[str, Any], name_prefix: str
             "apply": "tc-admin-apply-staging",
         },
         "attributes": {"integration": name_prefix},
+        "optimization": {
+            "integration-test": [
+                "taskcluster/fxci_config_taskgraph/**",
+                "taskcluster/kinds/firefoxci-artifact/kind.yml",
+                "taskcluster/kinds/integration-test/kind.yml",
+            ]
+        },
     }
     rewrite_docker_image(taskdesc)
     rewrite_private_fetches(taskdesc)

taskcluster/test/test_optimizations.py

@@ -0,0 +1,72 @@
+# 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/.
+
+import pytest
+from pytest_taskgraph import make_task
+
+from fxci_config_taskgraph import optimizations
+from fxci_config_taskgraph.util.constants import FIREFOXCI_ROOT_URL, STAGING_ROOT_URL
+
+
+@pytest.mark.parametrize(
+    "task,files_changed_should_remove,expected",
+    (
+        pytest.param(
+            {"provisionerId": "gecko-3", "workerType": "t-linux"},
+            True,
+            True,
+            id="removed",
+        ),
+        pytest.param(
+            {"provisionerId": "gecko-3", "workerType": "t-linux"},
+            False,
+            False,
+            id="kept via files changed",
+        ),
+        pytest.param(
+            {"provisionerId": "mozillavpn-3", "workerType": "b-linux-gcp-gw"},
+            True,
+            False,
+            id="kept via worker type",
+        ),
+    ),
+)
+def test_integration_test_strategy(
+    monkeypatch, mocker, task, files_changed_should_remove, expected
+):
+    monkeypatch.setenv("TASKCLUSTER_ROOT_URL", STAGING_ROOT_URL)
+
+    diff_output = """
+Registering resource: worker_pools
+! WorkerPool=code-analysis-3/linux-gw-gcp (changed: config)
+! WorkerPool=gecko-3/b-linux-2204-kvm-gcp (changed: config)
+! WorkerPool=mozillavpn-3/b-linux-gcp-gw (changed: config)
+"""
+
+    mock_proc = mocker.MagicMock()
+    mock_proc.stdout = diff_output
+
+    mock_run = mocker.patch.object(
+        optimizations.subprocess, "run", return_value=mock_proc
+    )
+
+    mock_files_changed = mocker.patch.object(
+        optimizations.SkipUnlessChanged,
+        "should_remove_task",
+        return_value=files_changed_should_remove,
+    )
+
+    task = make_task(label="foo", task_def=task)
+    opt = optimizations.IntegrationTestStrategy()
+    assert opt.should_remove_task(task, "foo", "bar") == expected
+
+    if mock_files_changed.called:
+        mock_files_changed.assert_called_once_with(task, "foo", "bar")
+
+    mock_run.assert_called_once()
+    args, kwargs = mock_run.call_args
+    cmd = " ".join(args[0])
+    assert "--environment firefoxci" in cmd
+    assert "--resources worker_pools" in cmd
+    assert kwargs["env"].get("TASKCLUSTER_ROOT_URL") == FIREFOXCI_ROOT_URL

taskcluster/test/test_transforms_integration_test.py

@@ -121,6 +121,13 @@ def test_basic(run_test):
         "dependencies": {"apply": "tc-admin-apply-staging"},
         "description": "test",
         "label": "gecko-foo",
+        "optimization": {
+            "integration-test": [
+                "taskcluster/fxci_config_taskgraph/**",
+                "taskcluster/kinds/firefoxci-artifact/kind.yml",
+                "taskcluster/kinds/integration-test/kind.yml",
+            ],
+        },
         "task": {
             "extra": {},
             "metadata": {"description": "test", "name": "gecko-foo"},