diff --git a/python/sites/python-test.txt b/python/sites/python-test.txt index c6cf70643dd7..58ad4f5efbce 100644 --- a/python/sites/python-test.txt +++ b/python/sites/python-test.txt @@ -1,4 +1,5 @@ pypi:pytest==7.0.1 +pypi:pytest-mock==3.12.0 pypi:Flask==2.1.3 # (indirect) avoids dependency on markupsafe >= 2.1.0, which is currently incompatible with glean-parser pypi:MarkupSafe==2.0.1 diff --git a/tools/tryselect/push.py b/tools/tryselect/push.py index 2b30c78e6aff..6b36cb42e194 100644 --- a/tools/tryselect/push.py +++ b/tools/tryselect/push.py @@ -97,8 +97,17 @@ def generate_try_task_config(method, labels, params=None, routes=None): # True). Their dependencies can be optimized though. params.setdefault("optimize_target_tasks", False) + # Remove selected labels from 'existing_tasks' parameter if present + if "existing_tasks" in params: + params["existing_tasks"] = { + label: tid + for label, tid in params["existing_tasks"].items() + if label not in labels + } + try_config = params.setdefault("try_task_config", {}) try_config.setdefault("env", {})["TRY_SELECTOR"] = method + try_config["tasks"] = sorted(labels) if routes: diff --git a/tools/tryselect/selectors/chooser/__init__.py b/tools/tryselect/selectors/chooser/__init__.py index 588d54302054..d6a32e08d032 100644 --- a/tools/tryselect/selectors/chooser/__init__.py +++ b/tools/tryselect/selectors/chooser/__init__.py @@ -31,6 +31,7 @@ class ChooserParser(BaseTryParser): "chemspill-prio", "disable-pgo", "env", + "existing-tasks", "gecko-profile", "path", "pernosco", diff --git a/tools/tryselect/selectors/fuzzy.py b/tools/tryselect/selectors/fuzzy.py index c4387ce2b1d4..e948100a0b01 100644 --- a/tools/tryselect/selectors/fuzzy.py +++ b/tools/tryselect/selectors/fuzzy.py @@ -114,6 +114,7 @@ class FuzzyParser(BaseTryParser): "chemspill-prio", "disable-pgo", "env", + "existing-tasks", "gecko-profile", "path", "pernosco", diff --git a/tools/tryselect/task_config.py b/tools/tryselect/task_config.py index ac1e244b71c3..c3124672a9a0 100644 --- a/tools/tryselect/task_config.py +++ b/tools/tryselect/task_config.py @@ -19,14 +19,17 @@ from contextlib import contextmanager from textwrap import dedent import mozpack.path as mozpath +import requests import six from mozbuild.base import BuildEnvironmentNotFoundException, MozbuildObject from mozversioncontrol import Repository +from taskgraph.util import taskcluster from .tasks import resolve_tests_by_suite +from .util.ssh import get_ssh_user -here = os.path.abspath(os.path.dirname(__file__)) -build = MozbuildObject.from_environment(cwd=here) +here = pathlib.Path(__file__).parent +build = MozbuildObject.from_environment(cwd=str(here)) @contextmanager @@ -155,13 +158,7 @@ class Pernosco(TryConfig): # log in. Prevent people with non-Mozilla addresses from using this # flag so they don't end up consuming time and resources only to # realize they can't actually log in and see the reports. - cmd = ["ssh", "-G", "hg.mozilla.org"] - output = subprocess.check_output( - cmd, universal_newlines=True - ).splitlines() - address = [ - l.rsplit(" ", 1)[-1] for l in output if l.startswith("user") - ][0] + address = get_ssh_user() if not address.endswith("@mozilla.com"): print( dedent( @@ -260,6 +257,70 @@ class Environment(TryConfig): } +class ExistingTasks(ParameterConfig): + TREEHERDER_PUSH_ENDPOINT = ( + "https://treeherder.mozilla.org/api/project/try/push/?count=1&author={user}" + ) + TREEHERDER_PUSH_URL = ( + "https://treeherder.mozilla.org/jobs?repo={branch}&revision={revision}" + ) + + arguments = [ + [ + ["-E", "--use-existing-tasks"], + { + "const": "last_try_push", + "default": None, + "nargs": "?", + "help": """ + Use existing tasks from a previous push. Without args this + uses your most recent try push. You may also specify + `rev=` where is the head revision of the + try push or `task-id=` where is the Decision + task id of the push. This last method even works for non-try + branches. + """, + }, + ] + ] + + def find_decision_task(self, use_existing_tasks): + branch = "try" + if use_existing_tasks == "last_try_push": + # Use existing tasks from user's previous try push. + user = get_ssh_user() + url = self.TREEHERDER_PUSH_ENDPOINT.format(user=user) + res = requests.get(url, headers={"User-Agent": "gecko-mach-try/1.0"}) + res.raise_for_status() + data = res.json() + if data["meta"]["count"] == 0: + raise Exception(f"Could not find a try push for '{user}'!") + revision = data["results"][0]["revision"] + + elif use_existing_tasks.startswith("rev="): + revision = use_existing_tasks[len("rev=") :] + + else: + raise Exception("Unable to parse '{use_existing_tasks}'!") + + url = self.TREEHERDER_PUSH_URL.format(branch=branch, revision=revision) + print(f"Using existing tasks from: {url}") + index_path = f"gecko.v2.{branch}.revision.{revision}.taskgraph.decision" + return taskcluster.find_task_id(index_path) + + def get_parameters(self, use_existing_tasks, **kwargs): + if not use_existing_tasks: + return + + if use_existing_tasks.startswith("task-id="): + tid = use_existing_tasks[len("task-id=") :] + else: + tid = self.find_decision_task(use_existing_tasks) + + label_to_task_id = taskcluster.get_artifact(tid, "public/label-to-taskid.json") + return {"existing_tasks": label_to_task_id} + + class RangeAction(Action): def __init__(self, min, max, *args, **kwargs): self.min = min @@ -418,7 +479,7 @@ class GeckoProfile(TryConfig): gecko_profile_entries, gecko_profile_features, gecko_profile_threads, - **kwargs + **kwargs, ): if profile or not all( s is None for s in (gecko_profile_features, gecko_profile_threads) @@ -548,6 +609,7 @@ all_task_configs = { "chemspill-prio": ChemspillPrio, "disable-pgo": DisablePgo, "env": Environment, + "existing-tasks": ExistingTasks, "gecko-profile": GeckoProfile, "path": Path, "pernosco": Pernosco, diff --git a/tools/tryselect/test/conftest.py b/tools/tryselect/test/conftest.py index ba27480a4064..d9cb7daee3eb 100644 --- a/tools/tryselect/test/conftest.py +++ b/tools/tryselect/test/conftest.py @@ -8,12 +8,19 @@ from unittest.mock import MagicMock import pytest import yaml from moztest.resolve import TestResolver +from responses import RequestsMock from taskgraph.graph import Graph from taskgraph.task import Task from taskgraph.taskgraph import TaskGraph from tryselect import push +@pytest.fixture +def responses(): + with RequestsMock() as rsps: + yield rsps + + @pytest.fixture def tg(request): if not hasattr(request.module, "TASKS"): diff --git a/tools/tryselect/test/python.toml b/tools/tryselect/test/python.toml index c624359b371c..f88156f69b6f 100644 --- a/tools/tryselect/test/python.toml +++ b/tools/tryselect/test/python.toml @@ -20,6 +20,8 @@ subsuite = "try" # shouldn't be run in parallel with those other tests. sequential = true +["test_push.py"] + ["test_release.py"] ["test_scriptworker.py"] diff --git a/tools/tryselect/test/test_push.py b/tools/tryselect/test/test_push.py new file mode 100644 index 000000000000..97f2e047d773 --- /dev/null +++ b/tools/tryselect/test/test_push.py @@ -0,0 +1,54 @@ +import mozunit +import pytest +from tryselect import push + + +@pytest.mark.parametrize( + "method,labels,params,routes,expected", + ( + pytest.param( + "fuzzy", + ["task-foo", "task-bar"], + None, + None, + { + "parameters": { + "optimize_target_tasks": False, + "try_task_config": { + "env": {"TRY_SELECTOR": "fuzzy"}, + "tasks": ["task-bar", "task-foo"], + }, + }, + "version": 2, + }, + id="basic", + ), + pytest.param( + "fuzzy", + ["task-foo"], + {"existing_tasks": {"task-foo": "123", "task-bar": "abc"}}, + None, + { + "parameters": { + "existing_tasks": {"task-bar": "abc"}, + "optimize_target_tasks": False, + "try_task_config": { + "env": {"TRY_SELECTOR": "fuzzy"}, + "tasks": ["task-foo"], + }, + }, + "version": 2, + }, + id="existing_tasks", + ), + ), +) +def test_generate_try_task_config(method, labels, params, routes, expected): + assert ( + push.generate_try_task_config(method, labels, params=params, routes=routes) + == expected + ) + + +if __name__ == "__main__": + mozunit.main() diff --git a/tools/tryselect/test/test_task_configs.py b/tools/tryselect/test/test_task_configs.py index fd646689fbfe..00adbf6228ce 100644 --- a/tools/tryselect/test/test_task_configs.py +++ b/tools/tryselect/test/test_task_configs.py @@ -3,7 +3,6 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. import inspect -import subprocess from argparse import ArgumentParser from textwrap import dedent @@ -11,6 +10,9 @@ import mozunit import pytest from tryselect.task_config import Pernosco, all_task_configs +TC_URL = "https://firefox-ci-tc.services.mozilla.com" +TH_URL = "https://treeherder.mozilla.org" + # task configs have a list of tests of the form (input, expected) TASK_CONFIG_TESTS = { "artifact": [ @@ -130,26 +132,25 @@ def test_task_configs(config_patch_resolver, task_config, args, expected): @pytest.fixture -def patch_pernosco_email_check(monkeypatch): - def inner(val): - def fake_check_output(*args, **kwargs): - return val - - monkeypatch.setattr(subprocess, "check_output", fake_check_output) +def patch_ssh_user(mocker): + def inner(user): + mock_stdout = mocker.Mock() + mock_stdout.stdout = dedent( + f""" + key1 foo + user {user} + key2 bar + """ + ) + return mocker.patch( + "tryselect.util.ssh.subprocess.run", return_value=mock_stdout + ) return inner -def test_pernosco(patch_pernosco_email_check): - patch_pernosco_email_check( - dedent( - """ - user foobar@mozilla.com - hostname hg.mozilla.com - """ - ) - ) - +def test_pernosco(patch_ssh_user): + patch_ssh_user("user@mozilla.com") parser = ArgumentParser() cfg = Pernosco() @@ -159,5 +160,94 @@ def test_pernosco(patch_pernosco_email_check): assert params == {"try_task_config": {"env": {"PERNOSCO": "1"}}} +def test_exisiting_tasks(responses, patch_ssh_user): + parser = ArgumentParser() + cfg = all_task_configs["existing-tasks"]() + cfg.add_arguments(parser) + + user = "user@example.com" + rev = "a" * 40 + task_id = "abc" + label_to_taskid = {"task-foo": "123", "task-bar": "456"} + + args = ["--use-existing-tasks"] + args = parser.parse_args(args) + + responses.add( + responses.GET, + f"{TH_URL}/api/project/try/push/?count=1&author={user}", + json={"meta": {"count": 1}, "results": [{"revision": rev}]}, + ) + + responses.add( + responses.GET, + f"{TC_URL}/api/index/v1/task/gecko.v2.try.revision.{rev}.taskgraph.decision", + json={"taskId": task_id}, + ) + + responses.add( + responses.GET, + f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json", + json=label_to_taskid, + ) + + m = patch_ssh_user(user) + params = cfg.get_parameters(**vars(args)) + assert params == {"existing_tasks": label_to_taskid} + + m.assert_called_once_with( + ["ssh", "-G", "hg.mozilla.org"], text=True, check=True, capture_output=True + ) + + +def test_exisiting_tasks_task_id(responses): + parser = ArgumentParser() + cfg = all_task_configs["existing-tasks"]() + cfg.add_arguments(parser) + + task_id = "abc" + label_to_taskid = {"task-foo": "123", "task-bar": "456"} + + args = ["--use-existing-tasks", f"task-id={task_id}"] + args = parser.parse_args(args) + + responses.add( + responses.GET, + f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json", + json=label_to_taskid, + ) + + params = cfg.get_parameters(**vars(args)) + assert params == {"existing_tasks": label_to_taskid} + + +def test_exisiting_tasks_rev(responses): + parser = ArgumentParser() + cfg = all_task_configs["existing-tasks"]() + cfg.add_arguments(parser) + + rev = "aaaaaa" + task_id = "abc" + label_to_taskid = {"task-foo": "123", "task-bar": "456"} + + args = ["--use-existing-tasks", f"rev={rev}"] + args = parser.parse_args(args) + + responses.add( + responses.GET, + f"{TC_URL}/api/index/v1/task/gecko.v2.try.revision.{rev}.taskgraph.decision", + json={"taskId": task_id}, + ) + + responses.add( + responses.GET, + f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json", + json=label_to_taskid, + ) + + params = cfg.get_parameters(**vars(args)) + assert params == {"existing_tasks": label_to_taskid} + + if __name__ == "__main__": mozunit.main() diff --git a/tools/tryselect/util/ssh.py b/tools/tryselect/util/ssh.py new file mode 100644 index 000000000000..7682306bc770 --- /dev/null +++ b/tools/tryselect/util/ssh.py @@ -0,0 +1,24 @@ +# 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 https://mozilla.org/MPL/2.0/. + +import subprocess + + +def get_ssh_user(host="hg.mozilla.org"): + ssh_config = subprocess.run( + ["ssh", "-G", host], + text=True, + check=True, + capture_output=True, + ).stdout + + lines = [l.strip() for l in ssh_config.splitlines()] + for line in lines: + if not line: + continue + key, value = line.split(" ", 1) + if key.lower() == "user": + return value + + raise Exception(f"Could not detect ssh user for '{host}'!")