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

smart_tests/commands/subset.py

@@ -188,6 +188,12 @@ def __init__(
                 "--print-input-snapshot-id",
                 help="Print the input snapshot ID returned from the server instead of the subset results"
             )] = False,
+            subset_id_file: Annotated[str | None, typer.Option(
+                "--subset-id-file",
+                help="Write the subset ID to a file",
+                metavar="FILE",
+                hidden=True
+            )] = None,
             bin_target: Annotated[Fraction | None, typer.Option(
                 "--bin",
                 help="Split subset into bins, e.g. --bin 1/4",
@@ -284,6 +290,7 @@ def warn(msg: str):
         self.prioritized_tests_mapping_file = prioritized_tests_mapping_file
         self.input_snapshot_id = input_snapshot_id.value if input_snapshot_id else None
         self.print_input_snapshot_id = print_input_snapshot_id
+        self.subset_id_file = subset_id_file
         self.bin_target = bin_target
         self.same_bin_files = list(same_bin_files)
         self.is_get_tests_from_guess = is_get_tests_from_guess
@@ -661,6 +668,18 @@ def _print_input_snapshot_id_value(self, subset_result: SubsetResult):
 
         click.echo(subset_result.subset_id)
 
+    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.",
+                self.tracking_client,
+                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 run(self):
         """called after tests are scanned to compute the optimized order"""
 
@@ -688,12 +707,17 @@ def run(self):
             warn_and_exit_if_fail_fast_mode("Error: no tests found matching the path.")
             if self.print_input_snapshot_id:
                 self._print_input_snapshot_id_value(subset_result)
+            if self.subset_id_file:
+                self._write_subset_id_to_file(subset_result)
             return
 
         if self.print_input_snapshot_id:
             self._print_input_snapshot_id_value(subset_result)
             return
 
+        if self.subset_id_file:
+            self._write_subset_id_to_file(subset_result)
+
         # TODO(Konboi): split subset isn't provided for smart-tests initial release
         # if split:
         #   click.echo("subset/{}".format(subset_result.subset_id))

tests/commands/test_subset.py

@@ -716,3 +716,275 @@ def test_subset_with_same_bin_file(self):
                     ],
                 ]],
             )
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_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,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            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", "file",
+            "--session", self.session,
+            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, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_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,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            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", "file",
+                "--session", self.session,
+                "--subset-id-file", id_file_path,
+                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(), f"{self.subsetting_id}\n")
+        finally:
+            os.unlink(id_file_path)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_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,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            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", "file",
+                "--session", self.session,
+                "--target", "30%",
+                "--subset-id-file", id_file_path,
+                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(), f"{self.subsetting_id}\n")
+        finally:
+            os.unlink(id_file_path)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_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,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            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", "file",
+                "--session", self.session,
+                "--rest", rest_file_path,
+                "--subset-id-file", id_file_path,
+                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(), f"{self.subsetting_id}\n")
+            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, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_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,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            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", "file",
+                "--session", self.session,
+                "--subset-id-file", id_file_path,
+                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)
+
+    @responses.activate
+    @mock.patch.dict(os.environ, {"SMART_TESTS_TOKEN": CliTestCase.smart_tests_token})
+    def test_subset_id_file_round_trip(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,
+            f"{get_base_url()}/intake/organizations/{self.organization}/workspaces/{self.workspace}/subset",
+            json=mock_json_response,
+            status=200,
+        )
+
+        with tempfile.NamedTemporaryFile(delete=False) as id_file:
+            id_file_path = id_file.name
+
+        try:
+            # Step 1: capture the subset ID into a file
+            result = self.cli(
+                "subset", "file",
+                "--session", self.session,
+                "--subset-id-file", id_file_path,
+                mix_stderr=False,
+                input=pipe,
+            )
+            self.assert_success(result)
+
+            # Step 2: feed the file back via --input-snapshot-id @file
+            result2 = self.cli(
+                "subset", "file",
+                "--session", self.session,
+                "--input-snapshot-id", f"@{id_file_path}",
+                mix_stderr=False,
+            )
+            self.assert_success(result2)
+            payload = self.decode_request_body(self.find_request('/subset', n=1).request.body)
+            self.assertEqual(payload.get('subsettingId'), self.subsetting_id)
+        finally:
+            os.unlink(id_file_path)