diff --git a/smart_tests/commands/subset.py b/smart_tests/commands/subset.py index 9f05b4aa9..4ee715ab7 100644 --- a/smart_tests/commands/subset.py +++ b/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)) diff --git a/tests/commands/test_subset.py b/tests/commands/test_subset.py index 8cd9fd5a5..48cf1d473 100644 --- a/tests/commands/test_subset.py +++ b/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)