Coverage for app / backend / src / tests / pytest_split / plugin.py: 61%
85 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-03 06:18 +0000
1# Vendored from https://github.com/jerry-git/pytest-split
2# Copyright (c) 2024 Jerry Pussinen
3#
4# Permission is hereby granted, free of charge, to any person obtaining
5# a copy of this software and associated documentation files (the
6# "Software"), to deal in the Software without restriction, including
7# without limitation the rights to use, copy, modify, merge, publish,
8# distribute, sublicense, and/or sell copies of the Software, and to
9# permit persons to whom the Software is furnished to do so, subject to
10# the following conditions:
11#
12# The above copyright notice and this permission notice shall be included
13# in all copies or substantial portions of the Software.
14#
15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
16# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
17# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
18# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
19# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
20# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
21# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
23from __future__ import annotations
25import json
26import os
27from typing import TYPE_CHECKING
29import pytest
30from _pytest.config import create_terminal_writer, hookimpl
31from _pytest.reports import TestReport
33from tests.pytest_split import algorithms
35if TYPE_CHECKING:
36 from _pytest import nodes
37 from _pytest.config import Config
38 from _pytest.config.argparsing import Parser
39 from _pytest.main import ExitCode
42# Ugly hack for freezegun compatibility: https://github.com/spulec/freezegun/issues/286
43STORE_DURATIONS_SETUP_AND_TEARDOWN_THRESHOLD = 60 * 10 # seconds
46def pytest_addoption(parser: Parser) -> None:
47 """
48 Declare pytest-split's options.
49 """
50 group = parser.getgroup(
51 "Split tests into groups which execution time is about the same. "
52 "Run with --store-durations to store information about test execution times."
53 )
54 group.addoption(
55 "--store-durations",
56 dest="store_durations",
57 action="store_true",
58 help="Store durations into '--durations-path'.",
59 )
60 group.addoption(
61 "--durations-path",
62 dest="durations_path",
63 help=(
64 "Path to the file in which durations are (to be) stored, "
65 "default is .test_durations in the current working directory"
66 ),
67 default=os.path.join(os.getcwd(), ".test_durations"),
68 )
69 group.addoption(
70 "--splits",
71 dest="splits",
72 type=int,
73 help="The number of groups to split the tests into",
74 )
75 group.addoption(
76 "--group",
77 dest="group",
78 type=int,
79 help="The group of tests that should be executed (first one is 1)",
80 )
81 group.addoption(
82 "--splitting-algorithm",
83 dest="splitting_algorithm",
84 type=str,
85 help=f"Algorithm used to split the tests. Choices: {algorithms.Algorithms.names()}",
86 default="least_duration",
87 choices=algorithms.Algorithms.names(),
88 )
89 group.addoption(
90 "--clean-durations",
91 dest="clean_durations",
92 action="store_true",
93 help=(
94 "Removes the test duration info for tests which are not present "
95 "while running the suite with '--store-durations'."
96 ),
97 )
100@pytest.hookimpl(tryfirst=True)
101def pytest_cmdline_main(config: Config) -> int | ExitCode | None:
102 """
103 Validate options.
104 """
105 group = config.getoption("group")
106 splits = config.getoption("splits")
108 if splits is None and group is None:
109 return None
111 if splits and group is None: 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true
112 raise pytest.UsageError("argument `--group` is required")
114 if group and splits is None: 114 ↛ 115line 114 didn't jump to line 115 because the condition on line 114 was never true
115 raise pytest.UsageError("argument `--splits` is required")
117 if splits < 1: 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true
118 raise pytest.UsageError("argument `--splits` must be >= 1")
120 if group < 1 or group > splits: 120 ↛ 121line 120 didn't jump to line 121 because the condition on line 120 was never true
121 raise pytest.UsageError(f"argument `--group` must be >= 1 and <= {splits}")
123 return None
126def pytest_configure(config: Config) -> None:
127 """
128 Enable the plugins we need.
129 """
130 if config.option.splits and config.option.group:
131 config.pluginmanager.register(PytestSplitPlugin(config), "pytestsplitplugin")
133 if config.option.store_durations:
134 config.pluginmanager.register(PytestSplitCachePlugin(config), "pytestsplitcacheplugin")
137class Base:
138 def __init__(self, config: Config) -> None:
139 """
140 Load durations and set up a terminal writer.
142 This logic is shared for both the split- and cache plugin.
143 """
144 self.config = config
145 self.writer = create_terminal_writer(self.config)
147 try:
148 with open(config.option.durations_path) as f:
149 self.cached_durations = json.loads(f.read())
150 except FileNotFoundError:
151 self.cached_durations = {}
153 # This code provides backwards compatibility after we switched
154 # from saving durations in a list-of-lists to a dict format
155 # Remove this when bumping to v1
156 if isinstance(self.cached_durations, list): 156 ↛ 157line 156 didn't jump to line 157 because the condition on line 156 was never true
157 self.cached_durations = dict(self.cached_durations)
160class PytestSplitPlugin(Base):
161 def __init__(self, config: Config):
162 super().__init__(config)
164 if not self.cached_durations: 164 ↛ 165line 164 didn't jump to line 165 because the condition on line 164 was never true
165 message = self.writer.markup(
166 "\n[pytest-split] No test durations found. Pytest-split will "
167 "split tests evenly when no durations are found. "
168 "\n[pytest-split] You can expect better results in consequent runs, "
169 "when test timings have been documented.\n"
170 )
171 self.writer.line(message)
173 @hookimpl(trylast=True)
174 def pytest_collection_modifyitems(self, config: Config, items: list[nodes.Item]) -> None:
175 """
176 Collect and select the tests we want to run, and deselect the rest.
177 """
178 splits: int = config.option.splits
179 group_idx: int = config.option.group
181 algo = algorithms.Algorithms[config.option.splitting_algorithm].value
182 groups = algo(splits, items, self.cached_durations)
183 group = groups[group_idx - 1]
185 items[:] = group.selected
186 config.hook.pytest_deselected(items=group.deselected)
188 self.writer.line(
189 self.writer.markup(
190 f"\n\n[pytest-split] Splitting tests with algorithm: {config.option.splitting_algorithm}"
191 )
192 )
193 self.writer.line(
194 self.writer.markup(
195 f"[pytest-split] Running group {group_idx}/{splits} (estimated duration: {group.duration:.2f}s)\n"
196 )
197 )
200class PytestSplitCachePlugin(Base):
201 """
202 The cache plugin writes durations to our durations file.
203 """
205 def pytest_sessionfinish(self) -> None:
206 """
207 Method is called by Pytest after the test-suite has run.
208 https://github.com/pytest-dev/pytest/blob/main/src/_pytest/main.py#L308
209 """
210 terminal_reporter = self.config.pluginmanager.get_plugin("terminalreporter")
211 test_durations: dict[str, float] = {}
213 for test_reports in terminal_reporter.stats.values(): # type: ignore[union-attr]
214 for test_report in test_reports:
215 if isinstance(test_report, TestReport):
216 # These ifs be removed after this is solved: # https://github.com/spulec/freezegun/issues/286
217 if test_report.duration < 0:
218 continue # pragma: no cover
219 if (
220 test_report.when in ("teardown", "setup")
221 and test_report.duration > STORE_DURATIONS_SETUP_AND_TEARDOWN_THRESHOLD
222 ):
223 # Ignore not legit teardown durations
224 continue # pragma: no cover
226 # Add test durations to map
227 if test_report.nodeid not in test_durations:
228 test_durations[test_report.nodeid] = 0
229 test_durations[test_report.nodeid] += test_report.duration
231 if self.config.option.clean_durations:
232 self.cached_durations = dict(test_durations)
233 else:
234 for k, v in test_durations.items():
235 self.cached_durations[k] = v
237 with open(self.config.option.durations_path, "w") as f:
238 json.dump(self.cached_durations, f, sort_keys=True, indent=4)
240 message = self.writer.markup(f"\n\n[pytest-split] Stored test durations in {self.config.option.durations_path}")
241 self.writer.line(message)