Introduce --clean-durations flag

src/pytest_split/plugin.py

@@ -63,6 +63,15 @@ def pytest_addoption(parser: "Parser") -> None:
         default="duration_based_chunks",
         choices=algorithms.Algorithms.names(),
     )
+    group.addoption(
+        "--clean-durations",
+        dest="clean_durations",
+        action="store_true",
+        help=(
+            "Removes the test duration info for tests which are not present "
+            "while running the suite with '--store-durations'."
+        ),
+    )
 
 
 @pytest.mark.tryfirst
@@ -198,11 +207,12 @@ def pytest_sessionfinish(self) -> None:
                         test_durations[test_report.nodeid] = 0
                     test_durations[test_report.nodeid] += test_report.duration
 
-        # Update the full cached-durations object
-        for k, v in test_durations.items():
-            self.cached_durations[k] = v
+        if self.config.option.clean_durations:
+            self.cached_durations = dict(test_durations)
+        else:
+            for k, v in test_durations.items():
+                self.cached_durations[k] = v
 
-        # Save durations
         with open(self.config.option.durations_path, "w") as f:
             json.dump(self.cached_durations, f)
 

tests/test_plugin.py

@@ -7,10 +7,12 @@
 
 pytest_plugins = ["pytester"]
 
+EXAMPLE_SUITE_TEST_COUNT = 10
+
 
 @pytest.fixture
 def example_suite(testdir):
-    testdir.makepyfile("".join(f"def test_{num}(): pass\n" for num in range(1, 11)))
+    testdir.makepyfile("".join(f"def test_{num}(): pass\n" for num in range(1, EXAMPLE_SUITE_TEST_COUNT + 1)))
     yield testdir
 
 
@@ -42,6 +44,50 @@ def test_it_stores(self, example_suite, durations_path):
         for duration in durations.values():
             assert isinstance(duration, float)
 
+    def test_it_overrides_existing_durations(self, example_suite, durations_path):
+        existing_duration_test_name = (
+            "test_it_overrides_existing_durations0/test_it_overrides_existing_durations.py::test_1"
+        )
+        old_value = 99
+        with open(durations_path, "w") as f:
+            json.dump({existing_duration_test_name: old_value}, f)
+
+        example_suite.runpytest("--store-durations", "--durations-path", durations_path)
+
+        with open(durations_path) as f:
+            durations = json.load(f)
+
+        assert durations[existing_duration_test_name] != old_value
+        assert len(durations) == EXAMPLE_SUITE_TEST_COUNT
+
+    def test_it_doesnt_remove_old_durations(self, example_suite, durations_path):
+        old_durations = {"test_old1": 1, "test_old2": 2}
+        with open(durations_path, "w") as f:
+            json.dump(old_durations, f)
+
+        example_suite.runpytest("--store-durations", "--durations-path", durations_path)
+
+        with open(durations_path) as f:
+            durations = json.load(f)
+
+        for item in old_durations:
+            assert item in durations.keys()
+        assert len(durations) == EXAMPLE_SUITE_TEST_COUNT + len(old_durations)
+
+    def test_it_removes_old_when_cli_flag_used(self, example_suite, durations_path):
+        old_durations = {"test_old1": 1, "test_old2": 2}
+        with open(durations_path, "w") as f:
+            json.dump(old_durations, f)
+
+        example_suite.runpytest("--store-durations", "--durations-path", durations_path, "--clean-durations")
+
+        with open(durations_path) as f:
+            durations = json.load(f)
+
+        for item in old_durations:
+            assert item not in durations.keys()
+        assert len(durations) == EXAMPLE_SUITE_TEST_COUNT
+
     def test_it_does_not_store_without_flag(self, example_suite, durations_path):
         example_suite.runpytest("--durations-path", durations_path)
         assert not os.path.exists(durations_path)