[AIENG-196] Add the new command for the early flake detection

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()

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')

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/<build-name>/test_sessions/<test-session-id>',
+    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)

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",

launchable/test_runners/file.py

@@ -52,3 +52,5 @@ def find_filename():
 
 
 split_subset = launchable.CommonSplitSubsetImpls(__name__).split_subset()
+
+launchable.CommonFlakeDetectionImpls(__name__).flake_detection()

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)

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

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('_', ' ')

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, "")