From 31b5f27b50a555dcb42a36c0b2926fe0acd33fb6 Mon Sep 17 00:00:00 2001 From: Naoto Ono Date: Wed, 10 Sep 2025 09:12:20 +0900 Subject: [PATCH 1/2] [AIENG-196] Add the new command for the early flake detection --- launchable/__main__.py | 2 + launchable/commands/retry/__init__.py | 13 +++ launchable/commands/retry/flake_detection.py | 86 ++++++++++++++++++++ launchable/test_runners/bazel.py | 2 + launchable/test_runners/file.py | 2 + launchable/test_runners/launchable.py | 30 +++++++ launchable/test_runners/raw.py | 2 + launchable/utils/commands.py | 1 + 8 files changed, 138 insertions(+) create mode 100644 launchable/commands/retry/__init__.py create mode 100644 launchable/commands/retry/flake_detection.py diff --git a/launchable/__main__.py b/launchable/__main__.py index ef81dfd9a..3128cde5c 100644 --- a/launchable/__main__.py +++ b/launchable/__main__.py @@ -12,6 +12,7 @@ from .commands.compare import compare from .commands.inspect import inspect from .commands.record import record +from .commands.retry import retry from .commands.split_subset import split_subset from .commands.stats import stats from .commands.subset import subset @@ -91,6 +92,7 @@ def main(ctx, log_level, plugin_dir, dry_run, skip_cert_verification): main.add_command(inspect) main.add_command(stats) main.add_command(compare) +main.add_command(retry) if __name__ == '__main__': main() diff --git a/launchable/commands/retry/__init__.py b/launchable/commands/retry/__init__.py new file mode 100644 index 000000000..7b11e6e7f --- /dev/null +++ b/launchable/commands/retry/__init__.py @@ -0,0 +1,13 @@ +import click + +from launchable.utils.click import GroupWithAlias + +from .flake_detection import flake_detection + + +@click.group(cls=GroupWithAlias) +def retry(): + pass + + +retry.add_command(flake_detection, 'flake-detection') diff --git a/launchable/commands/retry/flake_detection.py b/launchable/commands/retry/flake_detection.py new file mode 100644 index 000000000..8b225913c --- /dev/null +++ b/launchable/commands/retry/flake_detection.py @@ -0,0 +1,86 @@ +import os +import sys + +import click + +from launchable.app import Application +from launchable.commands.helper import find_or_create_session +from launchable.commands.test_path_writer import TestPathWriter +from launchable.utils.click import ignorable_error +from launchable.utils.env_keys import REPORT_ERROR_KEY +from launchable.utils.launchable_client import LaunchableClient +from launchable.utils.tracking import Tracking, TrackingClient + +from ...utils.commands import Command + + +@click.group(help="Early flake detection") +@click.option( + '--session', + 'session', + help='In the format builds//test_sessions/', + type=str, + required=True +) +@click.option( + '--confidence', + help='Confidence level for flake detection', + type=click.Choice(['low', 'medium', 'high'], case_sensitive=False), + required=True, +) +@click.pass_context +def flake_detection(ctx, confidence, session): + tracking_client = TrackingClient(Command.FLAKE_DETECTION, app=ctx.obj) + client = LaunchableClient(app=ctx.obj, tracking_client=tracking_client, test_runner=ctx.invoked_subcommand) + session_id = None + try: + session_id = find_or_create_session( + context=ctx, + session=session, + build_name=None, + tracking_client=tracking_client + ) + except click.UsageError as e: + click.echo(click.style(str(e), fg="red"), err=True) + sys.exit(1) + except Exception as e: + tracking_client.send_error_event( + event_name=Tracking.ErrorEvent.INTERNAL_CLI_ERROR, + stack_trace=str(e), + ) + if os.getenv(REPORT_ERROR_KEY): + raise e + else: + click.echo(ignorable_error(e), err=True) + if session_id is None: + return + + class FlakeDetection(TestPathWriter): + def __init__(self, app: Application): + super(FlakeDetection, self).__init__(app=app) + + def run(self): + test_paths = [] + try: + res = client.request( + "get", + "retry/flake-detection", + params={ + "confidence": confidence.upper(), + "session-id": os.path.basename(session_id), + "test-runner": ctx.invoked_subcommand}) + res.raise_for_status() + test_paths = res.json().get("testPaths", []) + if test_paths: + self.print(test_paths) + except Exception as e: + tracking_client.send_error_event( + event_name=Tracking.ErrorEvent.INTERNAL_CLI_ERROR, + stack_trace=str(e), + ) + if os.getenv(REPORT_ERROR_KEY): + raise e + else: + click.echo(ignorable_error(e), err=True) + + ctx.obj = FlakeDetection(app=ctx.obj) diff --git a/launchable/test_runners/bazel.py b/launchable/test_runners/bazel.py index 1904a926b..a76257d75 100644 --- a/launchable/test_runners/bazel.py +++ b/launchable/test_runners/bazel.py @@ -33,6 +33,8 @@ def subset(client): split_subset = launchable.CommonSplitSubsetImpls(__name__, formatter=lambda x: x[0]['name'] + ":" + x[1]['name']).split_subset() +launchable.CommonFlakeDetectionImpls(__name__, formatter=lambda x: x[0]['name'] + ":" + x[1]['name']).flake_detection() + @click.argument('workspace', required=True) @click.option('--build-event-json', 'build_event_json_files', help="set file path generated by --build_event_json_file", diff --git a/launchable/test_runners/file.py b/launchable/test_runners/file.py index 8c36746fc..dcfb4e2c5 100644 --- a/launchable/test_runners/file.py +++ b/launchable/test_runners/file.py @@ -52,3 +52,5 @@ def find_filename(): split_subset = launchable.CommonSplitSubsetImpls(__name__).split_subset() + +launchable.CommonFlakeDetectionImpls(__name__).flake_detection() diff --git a/launchable/test_runners/launchable.py b/launchable/test_runners/launchable.py index 92119326c..8635c2936 100644 --- a/launchable/test_runners/launchable.py +++ b/launchable/test_runners/launchable.py @@ -6,8 +6,10 @@ import click from launchable.commands.record.tests import tests as record_tests_cmd +from launchable.commands.retry.flake_detection import flake_detection as flake_detection_cmd from launchable.commands.split_subset import split_subset as split_subset_cmd from launchable.commands.subset import subset as subset_cmd +from launchable.testpath import unparse_test_path def cmdname(m): @@ -43,6 +45,10 @@ def subset(f): record.tests = lambda f: wrap(f, record_tests_cmd) +def flake_detection(f): + return wrap(f, flake_detection_cmd) + + def split_subset(f): return wrap(f, split_subset_cmd) @@ -165,3 +171,27 @@ def split_subset(client): client.run() return wrap(split_subset, split_subset_cmd, self.cmdname) + + +class CommonFlakeDetectionImpls: + def __init__( + self, + module_name, + formatter=unparse_test_path, + seperator="\n", + ): + self.cmdname = cmdname(module_name) + self._formatter = formatter + self._separator = seperator + + def flake_detection(self): + def flake_detection(client): + if self._formatter: + client.formatter = self._formatter + + if self._separator: + client.separator = self._separator + + client.run() + + return wrap(flake_detection, flake_detection_cmd, self.cmdname) diff --git a/launchable/test_runners/raw.py b/launchable/test_runners/raw.py index cf52d0245..4ed980769 100644 --- a/launchable/test_runners/raw.py +++ b/launchable/test_runners/raw.py @@ -47,6 +47,8 @@ def subset(client, test_path_file): split_subset = launchable.CommonSplitSubsetImpls(__name__, formatter=unparse_test_path, seperator='\n').split_subset() +launchable.CommonFlakeDetectionImpls(__name__).flake_detection() + @click.argument('test_result_files', required=True, type=click.Path(exists=True), nargs=-1) @launchable.record.tests diff --git a/launchable/utils/commands.py b/launchable/utils/commands.py index 6a8e357b6..bd65a52ba 100644 --- a/launchable/utils/commands.py +++ b/launchable/utils/commands.py @@ -8,6 +8,7 @@ class Command(Enum): RECORD_SESSION = 'RECORD_SESSION' SUBSET = 'SUBSET' COMMIT = 'COMMIT' + FLAKE_DETECTION = 'FLAKE_DETECTION' def display_name(self): return self.value.lower().replace('_', ' ') From 072ca6c6dfe733c42497143fe7bd18799fe4719b Mon Sep 17 00:00:00 2001 From: Naoto Ono Date: Wed, 10 Sep 2025 14:57:18 +0900 Subject: [PATCH 2/2] Add tests --- tests/commands/test_flake_detection.py | 83 ++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 tests/commands/test_flake_detection.py diff --git a/tests/commands/test_flake_detection.py b/tests/commands/test_flake_detection.py new file mode 100644 index 000000000..4589b8c1f --- /dev/null +++ b/tests/commands/test_flake_detection.py @@ -0,0 +1,83 @@ +import os +from unittest import mock + +import responses # type: ignore + +from launchable.utils.http_client import get_base_url +from tests.cli_test_case import CliTestCase + + +class FlakeDetectionTest(CliTestCase): + @responses.activate + @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token}) + def test_flake_detection_success(self): + mock_json_response = { + "testPaths": [ + [{"type": "file", "name": "test_flaky_1.py"}], + [{"type": "file", "name": "test_flaky_2.py"}], + ] + } + responses.add( + responses.GET, + f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/retry/flake-detection", + json=mock_json_response, + status=200, + ) + result = self.cli( + "retry", + "flake-detection", + "--session", + self.session, + "--confidence", + "high", + "file", + mix_stderr=False, + ) + self.assert_success(result) + self.assertIn("test_flaky_1.py", result.stdout) + self.assertIn("test_flaky_2.py", result.stdout) + + @responses.activate + @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token}) + def test_flake_detection_no_flakes(self): + mock_json_response = {"testPaths": []} + responses.add( + responses.GET, + f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/retry/flake-detection", + json=mock_json_response, + status=200, + ) + result = self.cli( + "retry", + "flake-detection", + "--session", + self.session, + "--confidence", + "low", + "file", + mix_stderr=False, + ) + self.assert_success(result) + self.assertEqual(result.stdout, "") + + @responses.activate + @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token}) + def test_flake_detection_api_error(self): + responses.add( + responses.GET, + f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/retry/flake-detection", + status=500, + ) + result = self.cli( + "retry", + "flake-detection", + "--session", + self.session, + "--confidence", + "medium", + "file", + mix_stderr=False, + ) + self.assert_exit_code(result, 0) + self.assertIn("Error", result.stderr) + self.assertEqual(result.stdout, "")