From 01eb010ecbde1682f3644b485db06915e3dedd28 Mon Sep 17 00:00:00 2001 From: Alexander Froch Date: Wed, 26 Aug 2026 15:31:55 +0200 Subject: [PATCH 1/3] Make the log level configurable and add timestamps to the log output The log level can now be set with the UPP_LOG_LEVEL environment variable or with the --log-level flag, which wins over the environment variable. Every log message carries the date and time, debug and info messages are written to stdout and warnings and above to stderr, so in a batch job the .err file only contains the problems. The console width is taken from the terminal instead of being hard coded, and the section banners adapt to it. --- changelog.md | 1 + docs/run.md | 16 ++ tests/unit/test_logger.py | 82 +++++++++- tests/unit/test_main.py | 6 + tests/unit/utils/test_check_input_samples.py | 14 +- upp/main.py | 27 ++-- upp/stages/hist.py | 5 +- upp/stages/merging.py | 5 +- upp/stages/normalisation.py | 5 +- upp/stages/resampling.py | 5 +- upp/utils/check_input_samples.py | 8 + upp/utils/list_components.py | 9 ++ upp/utils/logger.py | 161 ++++++++++++++++--- 13 files changed, 290 insertions(+), 54 deletions(-) diff --git a/changelog.md b/changelog.md index 9e75f405..2d076e25 100644 --- a/changelog.md +++ b/changelog.md @@ -2,6 +2,7 @@ ### [Latest] +- Make the log level configurable via `UPP_LOG_LEVEL` or `--log-level`, add timestamps to all log messages and write warnings and errors to stderr [#161](https://github.com/umami-hep/umami-preprocessing/pull/161) - Add container image documentation and generic Slurm/HTCondor submission scripts with config-driven component enumeration [#160](https://github.com/umami-hep/umami-preprocessing/pull/160) ### [v0.3.2](https://github.com/umami-hep/umami-preprocessing/releases/tag/v0.3.2) (04.08.2026) diff --git a/docs/run.md b/docs/run.md index bfb10376..9a151f97 100644 --- a/docs/run.md +++ b/docs/run.md @@ -15,6 +15,22 @@ For a comprehensive list of available flags, refer to `preprocess --help`. !!!info "If you are running on lxplus you may need to use `python3 upp/main.py` instead of `preprocess`" +### Logging + +By default UPP logs at the `INFO` level. You can change this with the `UPP_LOG_LEVEL` +environment variable or with the `--log-level` flag, which wins over the environment variable: + +```bash +export UPP_LOG_LEVEL=DEBUG +preprocess --config path/to/config.yaml --log-level WARNING +``` + +The available levels are `DEBUG`, `INFO`, `WARNING`, `ERROR` and `CRITICAL`. + +Every log message is prefixed with the date and time. Debug and info messages are written to +stdout while warnings, errors and critical messages are written to stderr, so in a batch job +the `.out` file holds the progress of the run and the `.err` file only the problems. + ### Splits The data is divided into three splits: training (`train`), validation (`val`), and testing (`test`). diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py index 45501d60..ff835030 100644 --- a/tests/unit/test_logger.py +++ b/tests/unit/test_logger.py @@ -1,12 +1,86 @@ from __future__ import annotations import logging +import re -from upp.utils.logger import setup_logger +import pytest +from upp.utils import logger as upp_logger +from upp.utils.logger import LOG_LEVEL_ENV, banner, resolve_log_level, setup_logger -def test_setup_logger(caplog): - caplog.set_level(logging.DEBUG) + +@pytest.fixture(autouse=True) +def reset_logger(monkeypatch): + monkeypatch.delenv(LOG_LEVEL_ENV, raising=False) + monkeypatch.setattr(upp_logger, "_configured_level", None) + + +def test_setup_logger(capsys): logger = setup_logger(level="DEBUG") logger.debug("Debug message") - assert "Debug message" in caplog.text + assert "Debug message" in capsys.readouterr().out + + +def test_resolve_log_level_default(): + assert resolve_log_level() == "INFO" + + +def test_resolve_log_level_from_env(monkeypatch): + monkeypatch.setenv(LOG_LEVEL_ENV, "debug") + assert resolve_log_level() == "DEBUG" + + +def test_resolve_log_level_argument_wins(monkeypatch): + monkeypatch.setenv(LOG_LEVEL_ENV, "DEBUG") + assert resolve_log_level("WARNING") == "WARNING" + + +def test_resolve_log_level_invalid(monkeypatch): + monkeypatch.setenv(LOG_LEVEL_ENV, "LOUD") + with pytest.raises(ValueError, match="Invalid log level LOUD"): + resolve_log_level() + + +def test_setup_logger_uses_env(monkeypatch): + monkeypatch.setenv(LOG_LEVEL_ENV, "DEBUG") + setup_logger() + assert logging.getLogger().level == logging.DEBUG + + +def test_setup_logger_keeps_first_level(monkeypatch): + setup_logger(level="DEBUG") + monkeypatch.setenv(LOG_LEVEL_ENV, "ERROR") + setup_logger() + assert logging.getLogger().level == logging.DEBUG + + +def test_records_are_split_between_stdout_and_stderr(capsys): + setup_logger(level="DEBUG") + logging.debug("a debug record") + logging.info("an info record") + logging.warning("a warning record") + logging.error("an error record") + + captured = capsys.readouterr() + assert "a debug record" in captured.out + assert "an info record" in captured.out + assert "a warning record" not in captured.out + assert "a warning record" in captured.err + assert "an error record" in captured.err + assert "an info record" not in captured.err + + +def test_records_have_a_timestamp(capsys): + setup_logger(level="INFO") + logging.info("a timestamped record") + captured = capsys.readouterr() + assert re.search(r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}", captured.out) + + +def test_banner_fits_the_console_width(capsys): + setup_logger(level="INFO") + logging.info(banner(" Title ")) + captured = capsys.readouterr() + lines = [line for line in captured.out.splitlines() if line.strip()] + assert len(lines) == 1 + assert " Title " in lines[0] diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 896652fa..6bd92db2 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -41,6 +41,7 @@ def test_parse_args_with_config(config_file): rw_merge_idx=None, files=None, skip_sample_check=False, + log_level=None, ) assert parsed_args == expected_args @@ -67,6 +68,7 @@ def test_parse_args_flags_not_given(config_file): rw_merge_idx=None, files=None, skip_sample_check=False, + log_level=None, ) assert parsed_args == expected_args @@ -104,6 +106,7 @@ def test_parse_args_flags_negative(config_file): rw_merge_idx=None, files=None, skip_sample_check=False, + log_level=None, ) assert parsed_args == expected_args @@ -139,6 +142,7 @@ def test_parse_args_flags_positive(config_file): rw_merge_idx=None, files=None, skip_sample_check=False, + log_level=None, ) assert parsed_args == expected_args @@ -176,6 +180,7 @@ def test_parse_args_component(config_file): rw_merge_idx=None, files=None, skip_sample_check=False, + log_level=None, ) assert parsed_args == expected_args @@ -213,6 +218,7 @@ def test_parse_args_region(config_file): rw_merge_idx=None, files=None, skip_sample_check=False, + log_level=None, ) assert parsed_args == expected_args diff --git a/tests/unit/utils/test_check_input_samples.py b/tests/unit/utils/test_check_input_samples.py index ebc6b9cd..28be0ae5 100644 --- a/tests/unit/utils/test_check_input_samples.py +++ b/tests/unit/utils/test_check_input_samples.py @@ -30,7 +30,7 @@ def info(self, *_a, **_k): def error(self, *_a, **_k): pass - monkeypatch.setattr(cis, "setup_logger", lambda: _Log()) + monkeypatch.setattr(cis, "setup_logger", lambda *_: _Log()) # H5Reader stub: we don't care about values here, just that it's called class _H5: @@ -74,7 +74,7 @@ def error(self, *_a, **_k): self._errors += 1 log = _Log() - monkeypatch.setattr(cis, "setup_logger", lambda: log) + monkeypatch.setattr(cis, "setup_logger", lambda *_: log) # H5Reader never called since pattern is invalid type class _H5: @@ -100,7 +100,9 @@ def __init__(self, ntuple_dir: Path): def test_main_calls_pipeline_with_parsed_args(monkeypatch, tmp_path): # Build fake args returned by parse_args - ns = Namespace(config_path=tmp_path / "cfg.yaml", deviation_factor=3.0, verbose=True) + ns = Namespace( + config_path=tmp_path / "cfg.yaml", deviation_factor=3.0, verbose=True, log_level=None + ) ns.config_path.write_text("") # so valid path conversion is happy if reached # Monkeypatch parse_args to return our namespace regardless of input @@ -155,7 +157,7 @@ def info(self, *_a, **_k): def error(self, *_a, **_k): pass - monkeypatch.setattr(cis, "setup_logger", lambda: _Log()) + monkeypatch.setattr(cis, "setup_logger", lambda *_: _Log()) # H5Reader stub class _H5: @@ -201,7 +203,7 @@ def info(self, *_a, **_k): def error(self, *_a, **_k): pass - monkeypatch.setattr(cis, "setup_logger", lambda: _Log()) + monkeypatch.setattr(cis, "setup_logger", lambda *_: _Log()) class _H5: def __init__(self, **_kwargs): @@ -276,7 +278,7 @@ def info(self, *_a, **_k): def error(self, *_a, **_k): pass - fake_logger.setup_logger = lambda: _Log() + fake_logger.setup_logger = lambda *_: _Log() # Inject fakes into sys.modules so the script can import them fakes = { diff --git a/upp/main.py b/upp/main.py index 98f45f5e..2febed98 100644 --- a/upp/main.py +++ b/upp/main.py @@ -30,7 +30,7 @@ from upp.stages.rw_merge import RWMerge from upp.stages.split_containers import SplitContainers from upp.utils.check_input_samples import run_input_sample_check -from upp.utils.logger import setup_logger +from upp.utils.logger import banner, setup_logger def parse_args(args: Any) -> argparse.Namespace: @@ -173,6 +173,13 @@ def parse_args(args: Any) -> argparse.Namespace: default=None, help="comma-separated list of files to use during the 'split-containers' stage ", ) + parser.add_argument( + "--log-level", + default=None, + type=str.upper, + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + help="Logging level. Overrides the UPP_LOG_LEVEL environment variable.", + ) args = parser.parse_args(args) d = vars(args) @@ -188,6 +195,7 @@ def parse_args(args: Any) -> argparse.Namespace: "reweight", "rw_merge", "rw_merge_idx", + "log_level", ] if not any(v for a, v in d.items() if a not in ignore): for v in d: @@ -204,7 +212,7 @@ def run_pp(args: argparse.Namespace) -> None: args : argparse.Namespace Parsed command line arguments """ - log = setup_logger() + log = setup_logger(args.log_level) # print start info log.info("[bold green]Starting preprocessing...") @@ -271,31 +279,28 @@ def run_pp(args: argparse.Namespace) -> None: # make plots if args.plot: - title = " Plotting " - log.info(f"[bold green]{title:-^100}") + log.info(banner(" Plotting ")) plot_resampling_dists(config=config, stage="initial") plot_resampling_dists(config=config, stage=args.split) # print end info end = datetime.now() - title = " Finished Preprocessing! " - log.info(f"[bold green]{title:-^100}") + log.info(banner(" Finished Preprocessing! ")) log.info(f"End time: {end.strftime('%Y-%m-%d %H:%M:%S')}") log.info(f"Elapsed time: {str(end - start).split('.')[0]}") def main(args: Any | None = None) -> None: args = parse_args(args) - log = setup_logger() + log = setup_logger(args.log_level) if args.split == "all": d = vars(args) for split in ["train", "val", "test"]: d["split"] = split - log.info(f"[bold blue]{'-' * 100}") - title = f" {args.split} " - log.info(f"[bold blue]{title:-^100}") - log.info(f"[bold blue]{'-' * 100}") + log.info(banner(style="bold blue")) + log.info(banner(f" {args.split} ", style="bold blue")) + log.info(banner(style="bold blue")) run_pp(args) else: run_pp(args) diff --git a/upp/stages/hist.py b/upp/stages/hist.py index 7ce8e9fc..96c4a438 100644 --- a/upp/stages/hist.py +++ b/upp/stages/hist.py @@ -12,7 +12,7 @@ from numpy.lib.recfunctions import structured_to_unstructured as s2u from scipy.stats import binned_statistic_dd -from upp.utils.logger import setup_logger +from upp.utils.logger import banner, setup_logger if TYPE_CHECKING: # pragma: no cover from upp.classes.preprocessing_config import PreprocessingConfig @@ -154,8 +154,7 @@ def create_histograms( return sampl_vars = config.sampl_cfg.vars - title = " Writing PDFs " - log.info(f"[bold green]{title:-^100}") + log.info(banner(" Writing PDFs ")) log.info( f"[bold green]Estimating PDFs using {config.num_global_objects_estimate_hist:,} objects..." ) diff --git a/upp/stages/merging.py b/upp/stages/merging.py index c5ac0433..0a210b24 100644 --- a/upp/stages/merging.py +++ b/upp/stages/merging.py @@ -10,7 +10,7 @@ import numpy as np from ftag.hdf5 import H5Writer, join_structured_arrays -from upp.utils.logger import ProgressBar +from upp.utils.logger import ProgressBar, banner from upp.utils.tools import path_append if TYPE_CHECKING: # pragma: no cover @@ -611,8 +611,7 @@ def write_components(self, sample: str | None, components: Components) -> None: def run(self): """Run merging of the components.""" - title = " Running Merging " - log.info(f"[bold green]{title:-^100}") + log.info(banner(" Running Merging ")) if not self.config.is_test or self.config.merge_test_samples: components = [(None, self.components)] diff --git a/upp/stages/normalisation.py b/upp/stages/normalisation.py index c2bc3617..69a0a1a7 100644 --- a/upp/stages/normalisation.py +++ b/upp/stages/normalisation.py @@ -8,7 +8,7 @@ import yaml from ftag.hdf5 import H5Reader -from upp.utils.logger import ProgressBar +from upp.utils.logger import ProgressBar, banner if TYPE_CHECKING: # pragma: no cover from upp.classes.preprocessing_config import PreprocessingConfig @@ -231,8 +231,7 @@ def write_class_dict(self, class_dict: dict) -> None: def run(self): """Run the normalisation calculation.""" - title = " Computing Normalisations " - log.info(f"[bold green]{title:-^100}") + log.info(banner(" Computing Normalisations ")) if self.config.rw_config is not None: fname = str(self.config.out_fname).replace(".h5", "_vds.h5") diff --git a/upp/stages/resampling.py b/upp/stages/resampling.py index 7180d958..f704c781 100644 --- a/upp/stages/resampling.py +++ b/upp/stages/resampling.py @@ -12,7 +12,7 @@ from upp.stages.hist import bin_global_objects from upp.stages.interpolation import subdivide_bins, upscale_array_regionally -from upp.utils.logger import ProgressBar +from upp.utils.logger import ProgressBar, banner if TYPE_CHECKING: # pragma: no cover from collections.abc import Generator @@ -497,8 +497,7 @@ def run(self, region: str | None = None, component: str | None = None): if component and not region: raise ValueError("Can't define component for resampling without region!") - title = " Running resampling " - log.info(f"[bold green]{title:-^100}") + log.info(banner(" Running resampling ")) log.info(f"Resampling method: {self.method or 'none'}") # Setup the different components and readers/writers and their sampling fraction diff --git a/upp/utils/check_input_samples.py b/upp/utils/check_input_samples.py index b60fa98e..60cffc4a 100644 --- a/upp/utils/check_input_samples.py +++ b/upp/utils/check_input_samples.py @@ -48,6 +48,13 @@ def parse_args(args: Any) -> argparse.Namespace: action="store_true", help="Print the final numbers to the terminal", ) + parser.add_argument( + "--log-level", + default=None, + type=str.upper, + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + help="Logging level. Overrides the UPP_LOG_LEVEL environment variable.", + ) args = parser.parse_args(args) return args @@ -233,6 +240,7 @@ def run_input_sample_check( def main(args: Any | None = None) -> None: args = parse_args(args) + setup_logger(args.log_level) # Load preprocessing config config = PreprocessingConfig.from_file( diff --git a/upp/utils/list_components.py b/upp/utils/list_components.py index 85b9cf05..16407342 100644 --- a/upp/utils/list_components.py +++ b/upp/utils/list_components.py @@ -8,6 +8,7 @@ from ftag.cli_utils import HelpFormatter, valid_path from upp.classes.preprocessing_config import PreprocessingConfig +from upp.utils.logger import setup_logger def parse_args(args: Any) -> argparse.Namespace: @@ -44,6 +45,13 @@ def parse_args(args: Any) -> argparse.Namespace: action="store_true", help="Only print the unique region names", ) + parser.add_argument( + "--log-level", + default=None, + type=str.upper, + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + help="Logging level. Overrides the UPP_LOG_LEVEL environment variable.", + ) return parser.parse_args(args) @@ -57,6 +65,7 @@ def main(args: Any | None = None) -> None: Command line arguments, by default None """ args = parse_args(args) + setup_logger(args.log_level) config = PreprocessingConfig.from_file( config_path=args.config, diff --git a/upp/utils/logger.py b/upp/utils/logger.py index 2f4f0f19..79aa1abf 100644 --- a/upp/utils/logger.py +++ b/upp/utils/logger.py @@ -1,9 +1,13 @@ from __future__ import annotations import logging +import os import sys +from datetime import datetime from functools import partial +from types import ModuleType +from ftag.utils.logging import get_log_level from rich.console import Console from rich.logging import RichHandler from rich.progress import ( @@ -14,18 +18,50 @@ TimeRemainingColumn, ) +# Name of the environment variable used to set the log level +LOG_LEVEL_ENV = "UPP_LOG_LEVEL" + +# Log level which is used when neither the command line nor the environment define one +DEFAULT_LOG_LEVEL = "INFO" + +# Timestamp which is prepended to every log message +LOG_TIME_FORMAT = "[%Y-%m-%d %H:%M:%S]" + +# Width of the time and level columns that Rich renders in front of every message, +# each including their trailing separator +_TIME_COL_WIDTH = len(datetime(2000, 1, 1).strftime(LOG_TIME_FORMAT)) + 1 +_LEVEL_COL_WIDTH = 9 + # Detect if the program is executed in an interactive terminal _IS_TTY = sys.stderr.isatty() +# Rich measures the terminal itself when no width is given. Under a batch system there is +# no terminal to measure, so fall back to a fixed width which COLUMNS can override. +_WIDTH = None if _IS_TTY else int(os.environ.get("COLUMNS", 120)) + # One console object is reused everywhere so that Rich keeps a consistent idea -# of whether it may emit ANSI control codes / animations. +# of whether it may emit ANSI control codes / animations. It is shared by the progress +# bar and the log records written to stdout (the .out file of a batch job), so that log +# messages are rendered above a running progress bar instead of on top of it. _console = Console( - width=100, + width=_WIDTH, force_terminal=_IS_TTY, force_interactive=_IS_TTY, no_color=not _IS_TTY, ) +# Console used for the log records which are written to stderr (the .err file of a batch job) +_stderr_console = Console( + stderr=True, + width=_WIDTH, + force_terminal=_IS_TTY, + force_interactive=_IS_TTY, + no_color=not _IS_TTY, +) + +# Level which was used in the last call of setup_logger() +_configured_level: str | None = None + # Template for the progress bar ProgressBar = partial( Progress, @@ -44,33 +80,116 @@ ) +def banner(title: str = "", style: str = "bold green") -> str: + """Build a horizontal rule with the title centred in the available log width. + + Parameters + ---------- + title : str, optional + Title to centre in the rule, by default "" + style : str, optional + Rich markup style applied to the rule, by default "bold green" + + Returns + ------- + str + Rule with Rich markup, ready to be passed to the logger + """ + width = max(_console.width - _TIME_COL_WIDTH - _LEVEL_COL_WIDTH, 20) + return f"[{style}]{title:-^{width}}" + + +def resolve_log_level(level: str | None = None) -> str: + """Resolve the log level from the command line, the environment or the default. + + Parameters + ---------- + level : str | None, optional + Log level given on the command line, which wins over the environment + variable, by default None + + Returns + ------- + str + Name of the resolved log level + + Raises + ------ + ValueError + If the resolved level is not a valid log level + """ + resolved = (level or os.environ.get(LOG_LEVEL_ENV) or DEFAULT_LOG_LEVEL).upper() + try: + get_log_level(resolved) + except ValueError as error: + raise ValueError(f"Invalid log level {resolved}") from error + return resolved + + +def _make_handler(console: Console) -> RichHandler: + """Create a Rich log handler which writes to the given console. + + Parameters + ---------- + console : Console + Console the handler writes to + + Returns + ------- + RichHandler + Handler with the timestamp column enabled + """ + return RichHandler( + console=console, + show_time=True, + show_path=False, + markup=True, + rich_tracebacks=True, + log_time_format=LOG_TIME_FORMAT, + omit_repeated_times=False, + ) + + # Helper for setup the logger -def setup_logger(level: str = "INFO"): +def setup_logger(level: str | None = None) -> ModuleType: """Set up the logger. Configure Rich logging so that colourful / interactive output is used when the program is attached to a terminal and plain text is written when it is executed under a batch system such as Slurm (where stdout / stderr are files). + Debug and info records are written to stdout, warnings and above to stderr. + + Parameters + ---------- + level : str | None, optional + Log level to use. If None, the level is taken from the UPP_LOG_LEVEL + environment variable and falls back to INFO. By default None + + Returns + ------- + ModuleType + The logging module, already configured """ + global _configured_level + + # Without an explicit level, keep the configuration of the first call + if level is None and _configured_level is not None: + return logging + FORMAT = "%(message)s" + resolved = resolve_log_level(level) - # In a batch job we create a console that never emits colour codes. - console = None - if not _IS_TTY: - console = Console( - width=120, - force_terminal=False, - force_interactive=False, - no_color=True, - ) - - handler = RichHandler( - show_time=False, - show_path=False, - markup=True, - rich_tracebacks=True, - console=console, - ) + stdout_handler = _make_handler(_console) + stdout_handler.addFilter(lambda record: record.levelno < logging.WARNING) - logging.basicConfig(level=level, format=FORMAT, handlers=[handler]) + stderr_handler = _make_handler(_stderr_console) + stderr_handler.setLevel(logging.WARNING) + + logging.basicConfig( + level=get_log_level(resolved), + format=FORMAT, + handlers=[stdout_handler, stderr_handler], + force=True, + ) + _configured_level = resolved return logging From 9cece7fd727010afbd448e0980f008635022cd60 Mon Sep 17 00:00:00 2001 From: Alexander Froch Date: Wed, 26 Aug 2026 15:33:33 +0200 Subject: [PATCH 2/3] Fix changelog link --- changelog.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog.md b/changelog.md index 2d076e25..3a13cc7f 100644 --- a/changelog.md +++ b/changelog.md @@ -2,7 +2,7 @@ ### [Latest] -- Make the log level configurable via `UPP_LOG_LEVEL` or `--log-level`, add timestamps to all log messages and write warnings and errors to stderr [#161](https://github.com/umami-hep/umami-preprocessing/pull/161) +- Make the log level configurable via `UPP_LOG_LEVEL` or `--log-level`, add timestamps to all log messages and write warnings and errors to stderr [#163](https://github.com/umami-hep/umami-preprocessing/pull/163) - Add container image documentation and generic Slurm/HTCondor submission scripts with config-driven component enumeration [#160](https://github.com/umami-hep/umami-preprocessing/pull/160) ### [v0.3.2](https://github.com/umami-hep/umami-preprocessing/releases/tag/v0.3.2) (04.08.2026) From 81d2d95d5e01a1024c6eb7e03e8591dd745ffa53 Mon Sep 17 00:00:00 2001 From: Alexander Froch Date: Wed, 26 Aug 2026 16:06:53 +0200 Subject: [PATCH 3/3] Add unit tests for the main entry point split handling --- tests/unit/test_main.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 6bd92db2..3bf10416 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4,7 +4,8 @@ from pytest import fixture -from upp.main import parse_args +import upp.main +from upp.main import main, parse_args @fixture @@ -222,3 +223,21 @@ def test_parse_args_region(config_file): ) assert parsed_args == expected_args + + +def test_main_runs_all_splits(config_file, monkeypatch): + splits = [] + monkeypatch.setattr(upp.main, "run_pp", lambda args: splits.append(args.split)) + + main(["--config", str(config_file), "--split", "all"]) + + assert splits == ["train", "val", "test"] + + +def test_main_runs_single_split(config_file, monkeypatch): + splits = [] + monkeypatch.setattr(upp.main, "run_pp", lambda args: splits.append(args.split)) + + main(["--config", str(config_file), "--split", "val"]) + + assert splits == ["val"]