From 0cbb911c857edf4188d7cf860ca7d6fb82e97297 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sondre=20Lilleb=C3=B8=20Gundersen?= <sondrelg@live.no>
Date: Mon, 31 May 2021 00:55:20 +0200
Subject: [PATCH] Add --store-durations back

---
 src/pytest_split/plugin.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/pytest_split/plugin.py b/src/pytest_split/plugin.py
index bc065ee..f470311 100644
--- a/src/pytest_split/plugin.py
+++ b/src/pytest_split/plugin.py
@@ -27,6 +27,12 @@ def pytest_addoption(parser: "Parser") -> None:
         "Run first the whole suite with --store-durations to save information "
         "about test execution times"
     )
+    group.addoption(
+        "--store-durations",
+        dest="store_durations",
+        action="store_true",
+        help="Store durations into '--durations-path'",
+    )
     group.addoption(
         "--splits",
         dest="splits",
@@ -67,9 +73,12 @@ def pytest_configure(config: "Config") -> None:
     elif config.option.splits and config.option.group:
         # Register plugin to run only if we received a splits and group arg
         config.pluginmanager.register(PytestSplitPlugin(config), "pytestsplitplugin")
+        config.pluginmanager.register(PytestSplitCachePlugin(config), "pytestsplitcacheplugin")
+    elif config.option.store_durations:
+        config.pluginmanager.register(PytestSplitCachePlugin(config), "pytestsplitcacheplugin")
 
 
-class PytestSplitPlugin:
+class PytestSplitBasePlugin:
     cache_file = "cache/pytest-split"
 
     def __init__(self, config: "Config") -> None:
@@ -93,6 +102,9 @@ def __init__(self, config: "Config") -> None:
                 "when test timings have been documented."
             )
 
+
+class PytestSplitPlugin(PytestSplitBasePlugin):
+
     @hookimpl(hookwrapper=True, tryfirst=True)
     def pytest_collection_modifyitems(self, config: "Config", items: "List[nodes.Item]") -> Generator[None, None, None]:
         """
@@ -122,10 +134,10 @@ def pytest_collection_modifyitems(self, config: "Config", items: "List[nodes.Ite
 
     @staticmethod
     def _split_tests(
-        splits: int,
-        group: int,
-        items: "List[nodes.Item]",
-        stored_durations: OrderedDict,
+            splits: int,
+            group: int,
+            items: "List[nodes.Item]",
+            stored_durations: OrderedDict,
     ) -> Tuple[int, int]:
         """
         Split tests by runtime.
@@ -201,6 +213,9 @@ def _split_tests(
 
         return selected, deselected
 
+
+class PytestSplitCachePlugin(PytestSplitBasePlugin):
+
     def pytest_sessionfinish(self) -> None:
         """
         Write test runtimes to cache.
@@ -219,8 +234,8 @@ def pytest_sessionfinish(self) -> None:
                     if test_report.duration < 0:
                         continue
                     if (
-                        getattr(test_report, "when", "") in ("teardown", "setup")
-                        and test_report.duration > STORE_DURATIONS_SETUP_AND_TEARDOWN_THRESHOLD
+                            getattr(test_report, "when", "") in ("teardown", "setup")
+                            and test_report.duration > STORE_DURATIONS_SETUP_AND_TEARDOWN_THRESHOLD
                     ):
                         # Ignore not legit teardown durations
                         continue