Add fallback-mode option for subset

smart_tests/commands/subset.py

@@ -2,6 +2,7 @@
 import json
 import os
 import pathlib
+import random
 import re
 import subprocess
 import sys
@@ -44,6 +45,12 @@ class SubsetUseCase(str, Enum):
     RECURRING = "recurring"
 
 
+class FallbackMode(str, Enum):
+    RUN_ALL = "run-all"
+    STOP = "stop"
+    RANDOM_SAMPLE = "random-sample"
+
+
 class SubsetResult:
     def __init__(
             self,
@@ -82,6 +89,14 @@ def from_test_paths(cls, test_paths: List[TestPath]) -> 'SubsetResult':
             is_observation=False
         )
 
+    @classmethod
+    def from_random_sample(cls, test_paths: List[TestPath], target: float) -> 'SubsetResult':
+        count = max(1, round(len(test_paths) * target))
+        sampled = random.sample(test_paths, min(count, len(test_paths)))
+        sampled_set = {id(t): t for t in sampled}
+        rest = [t for t in test_paths if id(t) not in sampled_set]
+        return cls(subset=sampled, rest=rest, subset_id='', summary={}, is_brainless=False, is_observation=False)
+
 
 # Where we take TestPath, we also accept a path name as a string.
 TestPathLike = str | TestPath
@@ -208,6 +223,14 @@ def __init__(
                 "--use-case",
                 hidden=True
             )] = None,
+            fallback_mode: Annotated[FallbackMode, typer.Option(
+                "--fallback-mode",
+                hidden=True,
+                help="Behavior when the subset API is unavailable or the model is untrained. "
+                     "'run-all' (default) runs all tests as usual; 'stop' exits with a non-zero status so CI halts; "
+                     "'random-sample' picks a random subset locally based on the count derived from --target "
+                     "(no duration estimates are available in this path).",
+            )] = FallbackMode.RUN_ALL,
             test_runner: Annotated[str | None, typer.Argument()] = None,
     ):
         super().__init__(app)
@@ -288,6 +311,7 @@ def warn(msg: str):
         self.same_bin_files = list(same_bin_files)
         self.is_get_tests_from_guess = is_get_tests_from_guess
         self.use_case = use_case
+        self.fallback_mode = fallback_mode
 
         self._validate_print_input_snapshot_option()
 
@@ -562,6 +586,23 @@ def _collect_potential_test_files(self):
         if not found:
             warn_and_exit_if_fail_fast_mode("Nothing that looks like a test file in the current git repository.")
 
+    def _fallback_result(self) -> SubsetResult:
+        if self.fallback_mode == FallbackMode.STOP:
+            click.echo(
+                "Warning: Smart Tests could not retrieve a subset. Stopping build (--fallback-mode=stop).",
+                err=True,
+            )
+            sys.exit(1)
+        elif self.fallback_mode == FallbackMode.RANDOM_SAMPLE:
+            target_fraction = float(self.target) if self.target is not None else 1.0
+            click.echo(
+                "Warning: Smart Tests could not retrieve a subset. Falling back to local random sample.",
+                err=True,
+            )
+            return SubsetResult.from_random_sample(self.test_paths, target_fraction)
+        else:
+            return SubsetResult.from_test_paths(self.test_paths)
+
     def request_subset(self) -> SubsetResult:
         # temporarily extend the timeout because subset API response has become slow
         # TODO: remove this line when API response return response
@@ -597,7 +638,7 @@ def request_subset(self) -> SubsetResult:
             )
             self.client.print_exception_and_recover(
                 e, "Warning: the service failed to subset. Falling back to running all tests")
-            return SubsetResult.from_test_paths(self.test_paths)
+            return self._fallback_result()
 
     def _requires_test_input(self) -> bool:
         return (
@@ -680,7 +721,7 @@ def run(self):
         if not self.session_id:
             # Session ID in --session is missing. It might be caused by
             # Launchable API errors.
-            subset_result = SubsetResult.from_test_paths(self.test_paths)
+            subset_result = self._fallback_result()
         else:
             subset_result = self.request_subset()
 
@@ -697,6 +738,11 @@ def run(self):
         # TODO(Konboi): split subset isn't provided for smart-tests initial release
         # if split:
         #   click.echo("subset/{}".format(subset_result.subset_id))
+        if subset_result.is_brainless:
+            click.echo("Your model is currently in training", err=True)
+            if self.fallback_mode != FallbackMode.RUN_ALL:
+                subset_result = self._fallback_result()
+
         output_subset, output_rests = subset_result.subset, subset_result.rest
 
         if subset_result.is_observation:
@@ -742,10 +788,6 @@ def run(self):
             ],
         ]
 
-        if subset_result.is_brainless:
-            click.echo(
-                "Your model is currently in training", err=True)
-
         click.echo(
             "Smart Tests created subset {} for build {} (test session {}) in workspace {}/{}".format(
                 subset_result.subset_id,

tests/commands/test_api_error.py

@@ -420,3 +420,104 @@ def assert_tracking_count(self, tracking, count: int):
                 if attempt > 10:
                     break
             self.assertEqual(tracking.call_count, count)
+
+
+class FallbackModeTest(CliTestCase):
+    test_files_dir = Path(__file__).parent.joinpath('../data/minitest/').resolve()
+
+    def _subset_args(self, rest_file_name, extra_args=()):
+        return (
+            "subset", "minitest",
+            "--target", "50%",
+            "--session", self.session,
+            "--rest", rest_file_name,
+            str(self.test_files_dir) + "/test/**/*.rb",
+        ) + tuple(extra_args)
+
+    # --- API error cases ---
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_api_error_fallback_stop(self):
+        responses.replace(
+            responses.POST,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            status=500)
+
+        with tempfile.NamedTemporaryFile(delete=False) as rest_file:
+            result = self.cli(*self._subset_args(rest_file.name, ("--fallback-mode", "stop")), mix_stderr=False)
+            self.assertEqual(result.exit_code, 1)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_api_error_fallback_random_sample(self):
+        responses.replace(
+            responses.POST,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            status=500)
+
+        with tempfile.NamedTemporaryFile(delete=False) as rest_file:
+            result = self.cli(*self._subset_args(rest_file.name, ("--fallback-mode", "random-sample")), mix_stderr=False)
+            self.assert_success(result)
+            all_tests = result.stdout.strip().split("\n") + Path(rest_file.name).read_text().strip().split("\n")
+            all_tests = [t for t in all_tests if t]
+            self.assertEqual(len(all_tests), 1)  # only one .rb test file in fixtures
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_api_error_fallback_run_all_default(self):
+        responses.replace(
+            responses.POST,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            status=500)
+
+        with tempfile.NamedTemporaryFile(delete=False) as rest_file:
+            result = self.cli(*self._subset_args(rest_file.name), mix_stderr=False)
+            self.assert_success(result)
+            self.assertIn("example_test.rb", result.stdout)
+
+    # --- Brainless mode cases ---
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_brainless_fallback_stop(self):
+        responses.replace(
+            responses.POST,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            json={"testPaths": [[{"type": "file", "name": "example_test.rb"}]],
+                  "rest": [], "subsettingId": 1, "isBrainless": True, "summary": {}},
+            status=200)
+
+        with tempfile.NamedTemporaryFile(delete=False) as rest_file:
+            result = self.cli(*self._subset_args(rest_file.name, ("--fallback-mode", "stop")), mix_stderr=False)
+            self.assertEqual(result.exit_code, 1)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_brainless_fallback_random_sample(self):
+        responses.replace(
+            responses.POST,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            json={"testPaths": [[{"type": "file", "name": "example_test.rb"}]],
+                  "rest": [], "subsettingId": 1, "isBrainless": True, "summary": {}},
+            status=200)
+
+        with tempfile.NamedTemporaryFile(delete=False) as rest_file:
+            result = self.cli(*self._subset_args(rest_file.name, ("--fallback-mode", "random-sample")), mix_stderr=False)
+            self.assert_success(result)
+            self.assertIn("example_test.rb", result.stdout)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_brainless_fallback_run_all_default(self):
+        responses.replace(
+            responses.POST,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            json={"testPaths": [[{"type": "file", "name": "example_test.rb"}]],
+                  "rest": [], "subsettingId": 1, "isBrainless": True, "summary": {}},
+            status=200)
+
+        with tempfile.NamedTemporaryFile(delete=False) as rest_file:
+            result = self.cli(*self._subset_args(rest_file.name), mix_stderr=False)
+            self.assert_success(result)
+            self.assertIn("example_test.rb", result.stdout)