[forward port] #1130, #1134

Konboi · Konboi · commit e7c72fc585af · 2025-09-22T12:30:53.000+09:00
#1130 #1134
diff --git a/smart_tests/__main__.py b/smart_tests/__main__.py
@@ -9,6 +9,7 @@
 import typer
 
 from smart_tests.app import Application
+from smart_tests.commands import detect_flakes
 from smart_tests.commands.record.tests import create_nested_commands as create_record_target_commands
 from smart_tests.commands.subset import create_nested_commands as create_subset_target_commands
 from smart_tests.utils.test_runner_registry import get_registry
@@ -150,6 +151,7 @@ def main(
     app.add_typer(inspect.app, name="inspect")
     app.add_typer(stats.app, name="stats")
     app.add_typer(compare.app, name="compare")
+    app.add_typer(detect_flakes.app, name="detect-flakes")
 
     # Add record-target as a sub-app to record command
     record.app.add_typer(record_target_app, name="test")  # Use NestedCommand version
@@ -162,6 +164,7 @@ def main(
     app.add_typer(verify.app, name="verify")
     app.add_typer(inspect.app, name="inspect")
     app.add_typer(stats.app, name="stats")
+    app.add_typer(detect_flakes.app, name="detect-flakes")
 
 app.callback()(main)
 
diff --git a/smart_tests/commands/detect_flakes.py b/smart_tests/commands/detect_flakes.py
@@ -0,0 +1,97 @@
+from enum import Enum
+import os
+from typing import Annotated
+import typer
+
+from smart_tests.app import Application
+from smart_tests.commands.test_path_writer import TestPathWriter
+from smart_tests.testpath import unparse_test_path
+from smart_tests.utils.commands import Command
+from smart_tests.utils.dynamic_commands import DynamicCommandBuilder, extract_callback_options
+from smart_tests.utils.env_keys import REPORT_ERROR_KEY
+from smart_tests.utils.exceptions import print_error_and_die
+from smart_tests.utils.session import get_session
+from smart_tests.utils.smart_tests_client import SmartTestsClient
+from smart_tests.utils.tracking import Tracking, TrackingClient
+from smart_tests.utils.typer_types import ignorable_error
+
+
+class DetectFlakesRetryThreshold(str, Enum):
+    LOW = "LOW"
+    MEDIUM = "MEDIUM"
+    HIGH = "HIGH"
+
+
+app = typer.Typer(name="detect-flakes", help="Detect flaky tests")
+
+
+@app.callback()
+def detect_flakes(
+    ctx: typer.Context,
+    session: Annotated[str, typer.Option(
+        "--session",
+        help="In the format builds/<build-name>/test_sessions/<test-session-id>",
+        metavar="SESSION"
+    )],
+    retry_threshold: Annotated[DetectFlakesRetryThreshold, typer.Option(
+        "--retry-threshold",
+        help="Thoroughness of how \"flake\" is detected",
+        case_sensitive=False,
+        show_default=True,
+    )] = DetectFlakesRetryThreshold.MEDIUM,
+):
+    app = ctx.obj
+    tracking_client = TrackingClient(Command.DETECT_FLAKE, app=app)
+    test_runner = getattr(ctx, 'test_runner', None)
+    client = SmartTestsClient(app=app, tracking_client=tracking_client, test_runner=test_runner)
+
+    test_session = None
+    try:
+        test_session = get_session(client=client, session=session)
+    except ValueError as e:
+        print_error_and_die(msg=str(e), tracking_client=tracking_client, event=Tracking.ErrorEvent.USER_ERROR)
+    except Exception as e:
+        if os.getenv(REPORT_ERROR_KEY):
+            raise e
+        else:
+            typer.echo(ignorable_error(e), err=True)
+
+    if test_session is None:
+        return
+
+    class FlakeDetection(TestPathWriter):
+        def __init__(self, app: Application):
+            super(FlakeDetection, self).__init__(app)
+
+        def run(self):
+            test_paths = []
+            try:
+                res = client.request(
+                    "get",
+                    "detect-flake",
+                    params={
+                        "confidence": retry_threshold.value.upper(),
+                        "session-id": os.path.basename(session),
+                        "test-runner": test_runner,
+                    })
+
+                res.raise_for_status()
+                test_paths = res.json().get("testPaths", [])
+                if test_paths:
+                    self.print(test_paths)
+                    typer.echo("Trying to retry the following tests:", err=True)
+                    for detail in res.json().get("testDetails", []):
+                        typer.echo(f"{detail.get('reason'): {unparse_test_path(detail.get('fullTestPath'))}}", err=True)
+            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:
+                    typer.echo(ignorable_error(e), err=True)
+
+    ctx.obj = FlakeDetection(app=ctx.obj)
+
+
diff --git a/smart_tests/test_runners/bazel.py b/smart_tests/test_runners/bazel.py
@@ -30,6 +30,9 @@ def subset(client):
     client.run()
 
 
+smart_tests.CommonDetectFlakesImpls(__name__).detect_flakes()
+
+
 @smart_tests.record.tests
 def record_tests(
     client,
diff --git a/smart_tests/test_runners/file.py b/smart_tests/test_runners/file.py
@@ -55,3 +55,6 @@ def find_filename():
     for r in reports:
         client.report(r)
     client.run()
+
+
+smart_tests.CommonDetectFlakesImpls(__name__).detect_flakes()
diff --git a/smart_tests/test_runners/rspec.py b/smart_tests/test_runners/rspec.py
@@ -2,3 +2,4 @@
 
 subset = smart_tests.CommonSubsetImpls(__name__).scan_files('*_spec.rb')
 record_tests = smart_tests.CommonRecordTestImpls(__name__).report_files()
+detect_flakes = smart_tests.CommonDetectFlakesImpls(__name__).detect_flakes()
diff --git a/smart_tests/test_runners/smart_tests.py b/smart_tests/test_runners/smart_tests.py
@@ -8,6 +8,7 @@
 
 from smart_tests.commands.record.tests import app as record_tests_cmd
 from smart_tests.commands.subset import app as subset_cmd
+from smart_tests.commands.detect_flakes import app as detect_flakes_cmd
 from smart_tests.utils.test_runner_registry import cmdname, create_test_runner_wrapper, get_registry
 
 
@@ -34,6 +35,13 @@ def subset(f):
     return f
 
 
+def detect_flakes(f):
+    test_runner_name = cmdname(f.__module__)
+    registry = get_registry()
+    registry.register_detect_flakes(test_runner_name, f)
+    return f
+
+
 record = types.SimpleNamespace()
 
 
@@ -152,3 +160,26 @@ def load_report_files(cls, client, source_roots, file_mask="*.xml"):
                 return
 
         client.run()
+
+
+class CommonDetectFlakesImpls:
+    def __init__(
+            self,
+            module_name,
+            formatter=None,
+            separator="\n",
+    ):
+        self.cmdname = cmdname(module_name)
+        self._formatter = formatter
+        self._separator = separator
+
+    def detect_flakes(self):
+        def detect_flakes(client):
+            if self._formatter:
+                client.formatter = self._formatter
+            if self._separator:
+                client.separator = self._separator
+
+            client.run()
+
+        return wrap(detect_flakes, detect_flakes_cmd, self.cmdname)
diff --git a/smart_tests/utils/commands.py b/smart_tests/utils/commands.py
@@ -8,6 +8,7 @@ class Command(Enum):
     RECORD_SESSION = 'RECORD_SESSION'
     SUBSET = 'SUBSET'
     COMMIT = 'COMMIT'
+    DETECT_FLAKE = 'DETECT_FLAKE'
 
     def display_name(self):
         return self.value.lower().replace('_', ' ')
diff --git a/smart_tests/utils/test_runner_registry.py b/smart_tests/utils/test_runner_registry.py
@@ -22,6 +22,7 @@ def __init__(self):
         self._subset_functions: Dict[str, Callable] = {}
         self._record_test_functions: Dict[str, Callable] = {}
         self._split_subset_functions: Dict[str, Callable] = {}
+        self._detect_flakes_functions: Dict[str, Callable] = {}
         # Callback to trigger when new test runners are registered
         self._on_register_callback: Callable[[], None] | None = None
 
@@ -47,6 +48,11 @@ def register_split_subset(self, test_runner_name: str, func: Callable) -> None:
         if self._on_register_callback:
             self._on_register_callback()
 
+    def register_detect_flakes(self, test_runner_name: str, func: Callable) -> None:
+        self._detect_flakes_functions[test_runner_name] = func
+        if self._on_register_callback:
+            self._on_register_callback()
+
     def get_subset_functions(self) -> Dict[str, Callable]:
         """Get all registered subset functions."""
         return self._subset_functions.copy()
diff --git a/tests/commands/test_detect_flakes.py b/tests/commands/test_detect_flakes.py
@@ -0,0 +1,99 @@
+import os
+from unittest import mock
+
+import responses  # type: ignore
+
+from smart_tests.utils.http_client import get_base_url
+from tests.cli_test_case import CliTestCase
+
+
+class DetectFlakeTest(CliTestCase):
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_detect_flakes_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}/detect-flake",
+            json=mock_json_response,
+            status=200,
+        )
+        result = self.cli(
+            "detect-flakes",
+            "file",
+            "--session", self.session,
+            "--retry-threshold", "high",
+            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, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_detect_flakes_without_retry_threshold_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}/detect-flake",
+            json=mock_json_response,
+            status=200,
+        )
+        result = self.cli(
+            "detect-flakes",
+            "file",
+            "--session", self.session,
+            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, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_detect_flakes_no_flakes(self):
+        mock_json_response = {"testPaths": []}
+        responses.add(
+            responses.GET,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/detect-flake",
+            json=mock_json_response,
+            status=200,
+        )
+        result = self.cli(
+            "detect-flakes",
+            "file",
+            "--session", self.session,
+            "--retry-threshold", "low",
+            mix_stderr=False,
+        )
+        self.assert_success(result)
+        self.assertEqual(result.stdout, "")
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_flake_detection_api_error(self):
+        responses.add(
+            responses.GET,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/detect-flake",
+            status=500,
+        )
+        result = self.cli(
+            "detect-flakes",
+            "file",
+            "--session", self.session,
+            "--retry-threshold", "medium",
+            mix_stderr=False,
+        )
+        self.assert_exit_code(result, 0)
+        self.assertIn("Error", result.stderr)
+        self.assertEqual(result.stdout, "")

Original file line number	Diff line number	Diff line change
`@@ -2,3 +2,4 @@`
`2`	`2`
`3`	`3`	`subset = smart_tests.CommonSubsetImpls(__name__).scan_files('*_spec.rb')`
`4`	`4`	`record_tests = smart_tests.CommonRecordTestImpls(__name__).report_files()`
	`5`	`+detect_flakes = smart_tests.CommonDetectFlakesImpls(__name__).detect_flakes()`