From da3de1bb0ad31cbc67bc54fdb0d9e283ad084edd Mon Sep 17 00:00:00 2001 From: Andrew Halberstadt Date: Fri, 13 Dec 2024 15:25:03 -0500 Subject: [PATCH] ci: implement an optimization for the integration tests --- taskcluster/fxci_config_taskgraph/__init__.py | 1 + .../fxci_config_taskgraph/optimizations.py | 45 ++++++++++++ .../transforms/integration_test.py | 7 ++ taskcluster/test/test_optimizations.py | 72 +++++++++++++++++++ .../test/test_transforms_integration_test.py | 7 ++ 5 files changed, 132 insertions(+) create mode 100644 taskcluster/fxci_config_taskgraph/optimizations.py create mode 100644 taskcluster/test/test_optimizations.py diff --git a/taskcluster/fxci_config_taskgraph/__init__.py b/taskcluster/fxci_config_taskgraph/__init__.py index 8796e00e..9face663 100644 --- a/taskcluster/fxci_config_taskgraph/__init__.py +++ b/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", ] ) diff --git a/taskcluster/fxci_config_taskgraph/optimizations.py b/taskcluster/fxci_config_taskgraph/optimizations.py new file mode 100644 index 00000000..e45c8a93 --- /dev/null +++ b/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) diff --git a/taskcluster/fxci_config_taskgraph/transforms/integration_test.py b/taskcluster/fxci_config_taskgraph/transforms/integration_test.py index 34f72146..761251b4 100644 --- a/taskcluster/fxci_config_taskgraph/transforms/integration_test.py +++ b/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) diff --git a/taskcluster/test/test_optimizations.py b/taskcluster/test/test_optimizations.py new file mode 100644 index 00000000..1214bc38 --- /dev/null +++ b/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 diff --git a/taskcluster/test/test_transforms_integration_test.py b/taskcluster/test/test_transforms_integration_test.py index 14eb66aa..1b12481a 100644 --- a/taskcluster/test/test_transforms_integration_test.py +++ b/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"},