Merge pull request #1328 from cloudbees-oss/backport-from-main-20260609

Add --subset-id-file option to capture subset ID without parsing stderr

Author: takanabe

launchable/commands/subset.py

@@ -223,6 +223,14 @@
     type=click.Choice(["one-commit", "feature-branch", "recurring"]),
     hidden=True,  # control PTS v2 test selection behavior. Non-committed, so hidden for now.
 )
+@click.option(
+    "--subset-id-file",
+    "subset_id_file",
+    help="Write the subset ID to a file",
+    metavar="FILE",
+    type=str,
+    hidden=True,
+)
 @click.pass_context
 def subset(
     context: click.core.Context,
@@ -253,6 +261,7 @@ def subset(
     is_get_tests_from_guess: bool = False,
     use_case: Optional[str] = None,
     similarity: Optional[float] = None,
+    subset_id_file: Optional[str] = None,
 ):
     app = context.obj
     tracking_client = TrackingClient(Command.SUBSET, app=app)
@@ -392,6 +401,7 @@ def __init__(self, app: Application):
             self.is_get_tests_from_guess = is_get_tests_from_guess
             self.is_output_exclusion_rules = is_output_exclusion_rules
             self.is_get_tests_from_guess = is_get_tests_from_guess
+            self.subset_id_file = subset_id_file
             super(Optimize, self).__init__(app=app)
 
         def _default_output_handler(self, output: List[TestPath], rests: List[TestPath]):
@@ -566,6 +576,17 @@ 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 _write_subset_id_to_file(self, subset_result: SubsetResult):
+            if not subset_result.subset_id:
+                print_error_and_die(
+                    "Subset request did not return a subset ID. Please re-run the command.",
+                    Tracking.ErrorEvent.INTERNAL_CLI_ERROR,
+                )
+
+            assert self.subset_id_file is not None  # Early type guard
+            with open(self.subset_id_file, 'w', encoding='utf-8') as f:
+                f.write(str(subset_result.subset_id) + '\n')
+
         def request_subset(self) -> SubsetResult:
             test_runner = context.invoked_subcommand
             # temporarily extend the timeout because subset API response has become slow
@@ -648,6 +669,9 @@ def run(self):
                 else:
                     self.output_handler(output_subset, output_rests)
 
+            if self.subset_id_file:
+                self._write_subset_id_to_file(subset_result)
+
             # When Launchable returns an error, the cli skips showing summary
             # report
             original_subset = subset_result.subset

tests/commands/test_subset.py

@@ -687,3 +687,233 @@ def test_subset_with_get_tests_from_guess(self):
         """
         payload = json.loads(gzip.decompress(responses.calls[1].request.body).decode())
         self.assertIn([{"type": "file", "name": "tests/commands/test_subset.py"}], payload.get("testPaths", []))
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token})
+    def test_subset_id_file_not_set(self):
+        pipe = "test_1.py\ntest_2.py\n"
+        mock_json_response = {
+            "testPaths": [
+                [{"type": "file", "name": "test_1.py"}],
+                [{"type": "file", "name": "test_2.py"}],
+            ],
+            "testRunner": "file",
+            "rest": [],
+            "subsettingId": self.subsetting_id,
+            "summary": {
+                "subset": {"duration": 10, "candidates": 2, "rate": 100},
+                "rest": {"duration": 0, "candidates": 0, "rate": 0},
+            },
+            "isObservation": False,
+        }
+        responses.replace(
+            responses.POST,
+            "{}/intake/organizations/{}/workspaces/{}/subset".format(
+                get_base_url(), self.organization, self.workspace),
+            json=mock_json_response,
+            status=200,
+        )
+
+        sentinel = tempfile.NamedTemporaryFile(delete=False)
+        sentinel_path = sentinel.name
+        sentinel.close()
+        os.unlink(sentinel_path)
+
+        result = self.cli(
+            "subset",
+            "--session", self.session,
+            "file",
+            mix_stderr=False,
+            input=pipe,
+        )
+        self.assert_success(result)
+        self.assertEqual(result.stdout, "test_1.py\ntest_2.py\n")
+        self.assertFalse(os.path.exists(sentinel_path))
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token})
+    def test_subset_id_file_written(self):
+        pipe = "test_1.py\ntest_2.py\n"
+        mock_json_response = {
+            "testPaths": [
+                [{"type": "file", "name": "test_1.py"}],
+                [{"type": "file", "name": "test_2.py"}],
+            ],
+            "testRunner": "file",
+            "rest": [],
+            "subsettingId": self.subsetting_id,
+            "summary": {
+                "subset": {"duration": 10, "candidates": 2, "rate": 100},
+                "rest": {"duration": 0, "candidates": 0, "rate": 0},
+            },
+            "isObservation": False,
+        }
+        responses.replace(
+            responses.POST,
+            "{}/intake/organizations/{}/workspaces/{}/subset".format(
+                get_base_url(), self.organization, self.workspace),
+            json=mock_json_response,
+            status=200,
+        )
+
+        with tempfile.NamedTemporaryFile(delete=False) as id_file:
+            id_file_path = id_file.name
+
+        try:
+            result = self.cli(
+                "subset",
+                "--session", self.session,
+                "--subset-id-file", id_file_path,
+                "file",
+                mix_stderr=False,
+                input=pipe,
+            )
+            self.assert_success(result)
+            self.assertEqual(result.stdout, "test_1.py\ntest_2.py\n")
+            with open(id_file_path) as f:
+                self.assertEqual(f.read(), "{}\n".format(self.subsetting_id))
+        finally:
+            os.unlink(id_file_path)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token})
+    def test_subset_id_file_with_target(self):
+        pipe = "test_1.py\ntest_2.py\ntest_3.py\n"
+        mock_json_response = {
+            "testPaths": [
+                [{"type": "file", "name": "test_1.py"}],
+            ],
+            "testRunner": "file",
+            "rest": [
+                [{"type": "file", "name": "test_2.py"}],
+                [{"type": "file", "name": "test_3.py"}],
+            ],
+            "subsettingId": self.subsetting_id,
+            "summary": {
+                "subset": {"duration": 5, "candidates": 1, "rate": 33},
+                "rest": {"duration": 10, "candidates": 2, "rate": 67},
+            },
+            "isObservation": False,
+        }
+        responses.replace(
+            responses.POST,
+            "{}/intake/organizations/{}/workspaces/{}/subset".format(
+                get_base_url(), self.organization, self.workspace),
+            json=mock_json_response,
+            status=200,
+        )
+
+        with tempfile.NamedTemporaryFile(delete=False) as id_file:
+            id_file_path = id_file.name
+
+        try:
+            result = self.cli(
+                "subset",
+                "--session", self.session,
+                "--target", "30%",
+                "--subset-id-file", id_file_path,
+                "file",
+                mix_stderr=False,
+                input=pipe,
+            )
+            self.assert_success(result)
+            self.assertEqual(result.stdout, "test_1.py\n")
+            with open(id_file_path) as f:
+                self.assertEqual(f.read(), "{}\n".format(self.subsetting_id))
+        finally:
+            os.unlink(id_file_path)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token})
+    def test_subset_id_file_with_rest(self):
+        pipe = "test_1.py\ntest_2.py\n"
+        mock_json_response = {
+            "testPaths": [
+                [{"type": "file", "name": "test_1.py"}],
+            ],
+            "testRunner": "file",
+            "rest": [
+                [{"type": "file", "name": "test_2.py"}],
+            ],
+            "subsettingId": self.subsetting_id,
+            "summary": {
+                "subset": {"duration": 5, "candidates": 1, "rate": 50},
+                "rest": {"duration": 5, "candidates": 1, "rate": 50},
+            },
+            "isObservation": False,
+        }
+        responses.replace(
+            responses.POST,
+            "{}/intake/organizations/{}/workspaces/{}/subset".format(
+                get_base_url(), self.organization, self.workspace),
+            json=mock_json_response,
+            status=200,
+        )
+
+        with tempfile.NamedTemporaryFile(delete=False) as id_file, \
+                tempfile.NamedTemporaryFile(delete=False) as rest_file:
+            id_file_path = id_file.name
+            rest_file_path = rest_file.name
+
+        try:
+            result = self.cli(
+                "subset",
+                "--session", self.session,
+                "--rest", rest_file_path,
+                "--subset-id-file", id_file_path,
+                "file",
+                mix_stderr=False,
+                input=pipe,
+            )
+            self.assert_success(result)
+            self.assertEqual(result.stdout, "test_1.py\n")
+            with open(id_file_path) as f:
+                self.assertEqual(f.read(), "{}\n".format(self.subsetting_id))
+            with open(rest_file_path) as f:
+                self.assertIn("test_2.py", f.read())
+        finally:
+            os.unlink(id_file_path)
+            os.unlink(rest_file_path)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token})
+    def test_subset_id_file_no_id_returned(self):
+        pipe = "test_1.py\n"
+        mock_json_response = {
+            "testPaths": [
+                [{"type": "file", "name": "test_1.py"}],
+            ],
+            "testRunner": "file",
+            "rest": [],
+            "subsettingId": "",
+            "summary": {
+                "subset": {"duration": 5, "candidates": 1, "rate": 100},
+                "rest": {"duration": 0, "candidates": 0, "rate": 0},
+            },
+            "isObservation": False,
+        }
+        responses.replace(
+            responses.POST,
+            "{}/intake/organizations/{}/workspaces/{}/subset".format(
+                get_base_url(), self.organization, self.workspace),
+            json=mock_json_response,
+            status=200,
+        )
+
+        with tempfile.NamedTemporaryFile(delete=False) as id_file:
+            id_file_path = id_file.name
+
+        try:
+            result = self.cli(
+                "subset",
+                "--session", self.session,
+                "--subset-id-file", id_file_path,
+                "file",
+                mix_stderr=False,
+                input=pipe,
+            )
+            self.assert_exit_code(result, 1)
+            self.assertIn("Subset request did not return a subset ID", result.stderr)
+        finally:
+            if os.path.exists(id_file_path):
+                os.unlink(id_file_path)