|
2 | 2 | from __future__ import annotations |
3 | 3 |
|
4 | 4 | import os |
| 5 | +import sys |
| 6 | +from dataclasses import dataclass |
| 7 | +from itertools import chain |
| 8 | +from typing import Iterator |
5 | 9 |
|
6 | 10 | import pytest |
7 | 11 |
|
| 12 | +if sys.version_info >= (3, 11): # pragma: >=3.11 cover |
| 13 | + import tomllib |
| 14 | +else: # pragma: <3.11 cover |
| 15 | + import tomli as tomllib |
| 16 | + |
8 | 17 |
|
9 | 18 | def pytest_addoption(parser: pytest.Parser) -> None: |
10 | 19 | """Add section to configuration files.""" |
11 | 20 | help_msg = "a line separated list of environment variables of the form (FLAG:)NAME=VALUE" |
12 | 21 | parser.addini("env", type="linelist", help=help_msg, default=[]) |
13 | 22 |
|
14 | 23 |
|
| 24 | +@dataclass |
| 25 | +class Entry: |
| 26 | + """Configuration entries.""" |
| 27 | + |
| 28 | + key: str |
| 29 | + value: str |
| 30 | + transform: bool |
| 31 | + skip_if_set: bool |
| 32 | + |
| 33 | + |
15 | 34 | @pytest.hookimpl(tryfirst=True) |
16 | 35 | def pytest_load_initial_conftests( |
17 | 36 | args: list[str], # noqa: ARG001 |
18 | 37 | early_config: pytest.Config, |
19 | 38 | parser: pytest.Parser, # noqa: ARG001 |
20 | 39 | ) -> None: |
21 | 40 | """Load environment variables from configuration files.""" |
22 | | - for line in early_config.getini("env"): |
23 | | - # INI lines e.g. D:R:NAME=VAL has two flags (R and D), NAME key, and VAL value |
24 | | - parts = line.partition("=") |
25 | | - ini_key_parts = parts[0].split(":") |
26 | | - flags = {k.strip().upper() for k in ini_key_parts[:-1]} |
27 | | - # R: is a way to designate whether to use raw value -> perform no transformation of the value |
28 | | - transform = "R" not in flags |
29 | | - # D: is a way to mark the value to be set only if it does not exist yet |
30 | | - skip_if_set = "D" in flags |
31 | | - key = ini_key_parts[-1].strip() |
32 | | - value = parts[2].strip() |
33 | | - |
34 | | - if skip_if_set and key in os.environ: |
| 41 | + for entry in _load_values(early_config): |
| 42 | + if entry.skip_if_set and entry.key in os.environ: |
35 | 43 | continue |
36 | 44 | # transformation -> replace environment variables, e.g. TEST_DIR={USER}/repo_test_dir. |
37 | | - os.environ[key] = value.format(**os.environ) if transform else value |
| 45 | + os.environ[entry.key] = entry.value.format(**os.environ) if entry.transform else entry.value |
| 46 | + |
| 47 | + |
| 48 | +def _load_values(early_config: pytest.Config) -> Iterator[Entry]: |
| 49 | + has_toml_conf = False |
| 50 | + for path in chain.from_iterable([[early_config.rootpath], early_config.rootpath.parents]): |
| 51 | + toml_file = path / "pyproject.toml" |
| 52 | + if toml_file.exists(): |
| 53 | + with toml_file.open("rb") as file_handler: |
| 54 | + config = tomllib.load(file_handler) |
| 55 | + if "tool" in config and "pytest_env" in config["tool"]: |
| 56 | + has_toml_conf = True |
| 57 | + for key, entry in config["tool"]["pytest_env"].get("env", {}).items(): |
| 58 | + if isinstance(entry, dict): |
| 59 | + value = str(entry["value"]) |
| 60 | + transform, skip_if_set = bool(entry.get("transform")), bool(entry.get("skip_if_set")) |
| 61 | + else: |
| 62 | + value, transform, skip_if_set = str(entry), False, False |
| 63 | + yield Entry(key, value, transform, skip_if_set) |
| 64 | + break |
| 65 | + |
| 66 | + if not has_toml_conf: |
| 67 | + for line in early_config.getini("env"): |
| 68 | + # INI lines e.g. D:R:NAME=VAL has two flags (R and D), NAME key, and VAL value |
| 69 | + parts = line.partition("=") |
| 70 | + ini_key_parts = parts[0].split(":") |
| 71 | + flags = {k.strip().upper() for k in ini_key_parts[:-1]} |
| 72 | + # R: is a way to designate whether to use raw value -> perform no transformation of the value |
| 73 | + transform = "R" not in flags |
| 74 | + # D: is a way to mark the value to be set only if it does not exist yet |
| 75 | + skip_if_set = "D" in flags |
| 76 | + key = ini_key_parts[-1].strip() |
| 77 | + value = parts[2].strip() |
| 78 | + yield Entry(key, value, transform, skip_if_set) |
0 commit comments