[update_lib] auto-mark original contents recovery (#6960)

scripts/update_lib/cmd_auto_mark.py

@@ -745,6 +745,7 @@ def auto_mark_file(
     # Strip reason-less markers so those tests fail normally and we capture
     # their error messages during the test run.
     contents = test_path.read_text(encoding="utf-8")
+    original_contents = contents
     contents, stripped_tests = strip_reasonless_expected_failures(contents)
     if stripped_tests:
         test_path.write_text(contents, encoding="utf-8")
@@ -761,11 +762,21 @@ def auto_mark_file(
         and not results.tests
         and not results.unexpected_successes
     ):
+        # Restore original contents before raising
+        if stripped_tests:
+            test_path.write_text(original_contents, encoding="utf-8")
         raise TestRunError(
             f"Test run failed for {test_name}. "
             f"Output: {results.stdout[-500:] if results.stdout else '(no output)'}"
         )
 
+    # If the run crashed (incomplete), restore original file so that markers
+    # for tests that never ran are preserved.  Only observed results will be
+    # re-applied below.
+    if not results.tests_result and stripped_tests:
+        test_path.write_text(original_contents, encoding="utf-8")
+        stripped_tests = set()
+
     contents = test_path.read_text(encoding="utf-8")
 
     all_failing_tests, unexpected_successes, error_messages = collect_test_changes(
@@ -863,11 +874,13 @@ def auto_mark_directory(
     # Strip reason-less markers from ALL files before running tests so those
     # tests fail normally and we capture their error messages.
     stripped_per_file: dict[pathlib.Path, set[tuple[str, str]]] = {}
+    original_per_file: dict[pathlib.Path, str] = {}
     for test_file in test_files:
         contents = test_file.read_text(encoding="utf-8")
-        contents, stripped = strip_reasonless_expected_failures(contents)
+        stripped_contents, stripped = strip_reasonless_expected_failures(contents)
         if stripped:
-            test_file.write_text(contents, encoding="utf-8")
+            original_per_file[test_file] = contents
+            test_file.write_text(stripped_contents, encoding="utf-8")
             stripped_per_file[test_file] = stripped
 
     test_name = get_test_module_name(test_dir)
@@ -882,11 +895,21 @@ def auto_mark_directory(
         and not results.tests
         and not results.unexpected_successes
     ):
+        # Restore original contents before raising
+        for fpath, original in original_per_file.items():
+            fpath.write_text(original, encoding="utf-8")
         raise TestRunError(
             f"Test run failed for {test_name}. "
             f"Output: {results.stdout[-500:] if results.stdout else '(no output)'}"
         )
 
+    # If the run crashed (incomplete), restore original files so that markers
+    # for tests that never ran are preserved.
+    if not results.tests_result and original_per_file:
+        for fpath, original in original_per_file.items():
+            fpath.write_text(original, encoding="utf-8")
+        stripped_per_file.clear()
+
     total_added = 0
     total_removed = 0
     total_regressions = 0

scripts/update_lib/tests/test_auto_mark.py

@@ -932,5 +932,154 @@ class TestA(unittest.TestCase):
                     auto_mark_directory(test_dir, verbose=False)
 
 
+class TestAutoMarkFileRestoresOnCrash(unittest.TestCase):
+    """Stripped markers must be restored when the test runner crashes."""
+
+    def test_stripped_markers_restored_when_crash(self):
+        """Markers stripped before run must be restored for unobserved tests on crash."""
+        test_code = f"""\
+import unittest
+
+class TestA(unittest.TestCase):
+    @unittest.expectedFailure  # {COMMENT}
+    def test_foo(self):
+        pass
+
+    @unittest.expectedFailure  # {COMMENT}
+    def test_bar(self):
+        pass
+
+    @unittest.expectedFailure  # {COMMENT}
+    def test_baz(self):
+        pass
+"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_file = pathlib.Path(tmpdir) / "test_example.py"
+            test_file.write_text(test_code)
+
+            # Simulate a crashed run that only observed test_foo (failed)
+            # test_bar and test_baz never ran due to crash
+            mock_result = TestResult()
+            mock_result.tests_result = ""  # no Tests result line (crash)
+            mock_result.tests = [
+                Test(
+                    name="test_foo",
+                    path="test.test_example.TestA.test_foo",
+                    result="fail",
+                    error_message="AssertionError: 1 != 2",
+                ),
+            ]
+
+            with mock.patch(
+                "update_lib.cmd_auto_mark.run_test", return_value=mock_result
+            ):
+                auto_mark_file(test_file, verbose=False)
+
+            contents = test_file.read_text()
+            # test_bar and test_baz were not observed — their markers must be restored
+            self.assertIn("def test_bar", contents)
+            self.assertIn("def test_baz", contents)
+            # Count expectedFailure markers: all 3 should be present
+            self.assertEqual(contents.count("expectedFailure"), 3, contents)
+
+    def test_stripped_markers_removed_when_complete_run(self):
+        """Markers are properly removed when the run completes normally."""
+        test_code = f"""\
+import unittest
+
+class TestA(unittest.TestCase):
+    @unittest.expectedFailure  # {COMMENT}
+    def test_foo(self):
+        pass
+
+    @unittest.expectedFailure  # {COMMENT}
+    def test_bar(self):
+        pass
+"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_file = pathlib.Path(tmpdir) / "test_example.py"
+            test_file.write_text(test_code)
+
+            # Simulate a complete run where test_foo fails but test_bar passes
+            mock_result = TestResult()
+            mock_result.tests_result = "FAILURE"  # normal completion
+            mock_result.tests = [
+                Test(
+                    name="test_foo",
+                    path="test.test_example.TestA.test_foo",
+                    result="fail",
+                    error_message="AssertionError",
+                ),
+            ]
+            # test_bar passes → shows as unexpected success
+            mock_result.unexpected_successes = [
+                Test(
+                    name="test_bar",
+                    path="test.test_example.TestA.test_bar",
+                    result="unexpected success",
+                ),
+            ]
+
+            with mock.patch(
+                "update_lib.cmd_auto_mark.run_test", return_value=mock_result
+            ):
+                auto_mark_file(test_file, verbose=False)
+
+            contents = test_file.read_text()
+            # test_foo should still have marker (re-added)
+            self.assertEqual(contents.count("expectedFailure"), 1, contents)
+            self.assertIn("def test_foo", contents)
+
+
+class TestAutoMarkDirectoryRestoresOnCrash(unittest.TestCase):
+    """Stripped markers must be restored for directory runs that crash."""
+
+    def test_stripped_markers_restored_when_crash(self):
+        test_code = f"""\
+import unittest
+
+class TestA(unittest.TestCase):
+    @unittest.expectedFailure  # {COMMENT}
+    def test_foo(self):
+        pass
+
+    @unittest.expectedFailure  # {COMMENT}
+    def test_bar(self):
+        pass
+"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_dir = pathlib.Path(tmpdir) / "test_example"
+            test_dir.mkdir()
+            test_file = test_dir / "test_sub.py"
+            test_file.write_text(test_code)
+
+            mock_result = TestResult()
+            mock_result.tests_result = ""  # crash
+            mock_result.tests = [
+                Test(
+                    name="test_foo",
+                    path="test.test_example.test_sub.TestA.test_foo",
+                    result="fail",
+                ),
+            ]
+
+            with (
+                mock.patch(
+                    "update_lib.cmd_auto_mark.run_test", return_value=mock_result
+                ),
+                mock.patch(
+                    "update_lib.cmd_auto_mark.get_test_module_name",
+                    side_effect=lambda p: (
+                        "test_example" if p == test_dir else "test_example.test_sub"
+                    ),
+                ),
+            ):
+                auto_mark_directory(test_dir, verbose=False)
+
+            contents = test_file.read_text()
+            # Both markers must be present (unobserved test_bar restored)
+            self.assertEqual(contents.count("expectedFailure"), 2, contents)
+
+
 if __name__ == "__main__":
     unittest.main()