diff --git a/src/duct/__main__.py b/src/duct/__main__.py index ea264d8d..75e12737 100644 --- a/src/duct/__main__.py +++ b/src/duct/__main__.py @@ -2,10 +2,9 @@ from __future__ import annotations import argparse from collections.abc import Iterable -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass, field, fields from datetime import datetime from enum import Enum -import glob import json import os from pathlib import Path @@ -14,7 +13,7 @@ import sys import threading import time -from typing import IO, Any, TextIO +from typing import IO, Any, Generator, TextIO from . import __version__ ENV_PREFIXES = ("PBS_", "SLURM_", "OSG") @@ -92,6 +91,18 @@ def max(self, other: ProcessStats) -> ProcessStats: ) +@dataclass +class Suffix: + stdout: str = "stdout" + stderr: str = "stderr" + usage: str = "usage.json" + info: str = "info.json" + + def __iter__(self) -> Generator: + for each in fields(self): + yield getattr(self, each.name) + + @dataclass class Sample: stats: dict[int, ProcessStats] = field(default_factory=dict) @@ -230,7 +241,7 @@ def write_pid_samples(self) -> None: assert self._sample is not None # First time only if self._resource_stats_log_path is None: - self._resource_stats_log_path = f"{self.output_prefix}usage.json" + self._resource_stats_log_path = f"{self.output_prefix}{Suffix.usage}" clobber_or_clear(self._resource_stats_log_path, self.clobber) with open(self._resource_stats_log_path, "a") as resource_statistics_log: @@ -467,7 +478,7 @@ def prepare_outputs( stderr: TextIO | TailPipe | int | None if capture_outputs.has_stdout(): - stdout_path = f"{output_prefix}stdout" + stdout_path = f"{output_prefix}{Suffix.stdout}" clobber_or_clear(stdout_path, clobber) Path(stdout_path).touch() # File must exist for TailPipe to read if outputs.has_stdout(): @@ -481,7 +492,7 @@ def prepare_outputs( stdout = subprocess.DEVNULL if capture_outputs.has_stderr(): - stderr_path = f"{output_prefix}stderr" + stderr_path = f"{output_prefix}{Suffix.stderr}" clobber_or_clear(stderr_path, clobber) Path(stderr_path).touch() if outputs.has_stderr(): @@ -504,23 +515,22 @@ def safe_close_files(file_list: Iterable[Any]) -> None: pass -def ensure_directories(path: str, clobber: bool = False) -> None: - # Enforcing no path* files is perhaps overzealous, but this helps prevent - # conflicts between versions, prevents probable user errors, and leaves - # room for plugins down the road. - possible_conflicts = glob.glob(path + "*") - if possible_conflicts and not clobber: +def ensure_directories(prefix: str, clobber: bool = False) -> None: + conflicts = [ + f"{prefix}{suffix}" for suffix in Suffix() if Path(f"{prefix}{suffix}").exists() + ] + if conflicts and not clobber: raise FileExistsError( - "Possibly conflicting files:\n" - + "\n".join(f"- {fp}" for fp in possible_conflicts) + "Conflicting files:\n" + + "\n".join(f"- {path}" for path in conflicts) + "\nUse --clobber to overrwrite conflicting files." ) - if path.endswith(os.sep): # If it ends in "/" (for linux) treat as a dir - os.makedirs(path, exist_ok=True) + if prefix.endswith(os.sep): # If it ends in "/" (for linux) treat as a dir + os.makedirs(prefix, exist_ok=True) else: # Path does not end with a separator, treat the last part as a filename - directory = os.path.dirname(path) + directory = os.path.dirname(prefix) if directory: os.makedirs(directory, exist_ok=True) @@ -594,7 +604,7 @@ def execute(args: Arguments) -> None: if args.record_types.has_system_summary(): report.collect_environment() report.get_system_info() - system_info_path = f"{args.output_prefix}info.json".format( + system_info_path = f"{args.output_prefix}{Suffix.info}".format( pid=duct_pid, datetime_filesafe=datetime_filesafe ) clobber_or_clear(system_info_path, clobber=args.clobber) diff --git a/test/test_execution.py b/test/test_execution.py index 55ecd35b..c4aeaef1 100644 --- a/test/test_execution.py +++ b/test/test_execution.py @@ -2,7 +2,7 @@ from pathlib import Path from unittest import mock from utils import assert_files -from duct.__main__ import Arguments, Outputs, RecordTypes, execute +from duct.__main__ import Arguments, Outputs, RecordTypes, Suffix, execute TEST_SCRIPT = str(Path(__file__).with_name("data") / "test_script.py") @@ -21,7 +21,7 @@ def test_sanity_green(temp_output_dir: str) -> None: ) execute(args) # When runtime < sample_interval, we won't have a usage.json - expected_files = ["stdout", "stderr", "info.json"] + expected_files = [Suffix.stdout, Suffix.stderr, Suffix.info] assert_files(temp_output_dir, expected_files, exists=True) @@ -42,10 +42,10 @@ def test_sanity_red(temp_output_dir: str) -> None: mock_stdout.write.assert_has_calls([mock.call("Exit Code: 1")]) # We still should execute normally - expected_files = ["stdout", "stderr", "info.json"] + expected_files = [Suffix.stdout, Suffix.stderr, Suffix.info] assert_files(temp_output_dir, expected_files, exists=True) # But no polling of the already failed command - not_expected_files = ["usage.json"] + not_expected_files = [Suffix.usage] assert_files(temp_output_dir, not_expected_files, exists=False) @@ -62,7 +62,7 @@ def test_outputs_full(temp_output_dir: str) -> None: clobber=False, ) execute(args) - expected_files = ["stdout", "stderr", "info.json", "usage.json"] + expected_files = [Suffix.stdout, Suffix.stderr, Suffix.info, Suffix.usage] assert_files(temp_output_dir, expected_files, exists=True) @@ -79,9 +79,9 @@ def test_outputs_passthrough(temp_output_dir: str) -> None: clobber=False, ) execute(args) - expected_files = ["info.json", "usage.json"] + expected_files = [Suffix.info, Suffix.usage] assert_files(temp_output_dir, expected_files, exists=True) - not_expected_files = ["stdout", "stderr"] + not_expected_files = [Suffix.stdout, Suffix.stderr] assert_files(temp_output_dir, not_expected_files, exists=False) @@ -100,7 +100,7 @@ def test_outputs_capture(temp_output_dir: str) -> None: execute(args) # TODO make this work assert mock.call("this is of test of STDOUT\n") not in mock_stdout.write.mock_calls - expected_files = ["stdout", "stderr", "info.json", "usage.json"] + expected_files = [Suffix.stdout, Suffix.stderr, Suffix.info, Suffix.usage] assert_files(temp_output_dir, expected_files, exists=True) @@ -119,10 +119,10 @@ def test_outputs_none(temp_output_dir: str) -> None: execute(args) # assert mock.call("this is of test of STDOUT\n") not in mock_stdout.write.mock_calls - expected_files = ["info.json", "usage.json"] + expected_files = [Suffix.info, Suffix.usage] assert_files(temp_output_dir, expected_files, exists=True) - not_expected_files = ["stdout", "stderr"] + not_expected_files = [Suffix.stdout, Suffix.stderr] assert_files(temp_output_dir, not_expected_files, exists=False) @@ -139,9 +139,9 @@ def test_exit_before_first_sample(temp_output_dir: str) -> None: clobber=False, ) execute(args) - expected_files = ["stdout", "stderr", "info.json"] + expected_files = [Suffix.stdout, Suffix.stderr, Suffix.info] assert_files(temp_output_dir, expected_files, exists=True) - not_expected_files = ["usage.json"] + not_expected_files = [Suffix.usage] assert_files(temp_output_dir, not_expected_files, exists=False) @@ -159,5 +159,5 @@ def test_run_less_than_report_interval(temp_output_dir: str) -> None: ) execute(args) # Specifically we need to assert that usage.json gets written anyway. - expected_files = ["stdout", "stderr", "usage.json", "info.json"] + expected_files = [Suffix.stdout, Suffix.stderr, Suffix.usage, Suffix.info] assert_files(temp_output_dir, expected_files, exists=True) diff --git a/test/test_helpers.py b/test/test_helpers.py index 6db9310a..9fce3db2 100644 --- a/test/test_helpers.py +++ b/test/test_helpers.py @@ -1,4 +1,5 @@ from __future__ import annotations +import pathlib from unittest import mock import pytest from duct.__main__ import clobber_or_clear, ensure_directories @@ -18,27 +19,26 @@ def test_ensure_directories_with_dirs(mock_mkdir: mock.MagicMock, path: str) -> mock_mkdir.assert_called_once_with(path, exist_ok=True) +@mock.patch("duct.__main__.Path", spec=pathlib.Path) @mock.patch("duct.__main__.os.makedirs") -@mock.patch("duct.__main__.glob.glob") def test_ensure_directories_with_conflicts( - mock_glob: mock.MagicMock, mock_mkdir: mock.MagicMock + mock_mkdir: mock.MagicMock, mock_path: mock.MagicMock ) -> None: - mock_path = "mock_path" - mock_glob.return_value = ["possibly", "conflicting", "files"] + mock_path.return_value.exists.return_value = True with pytest.raises(FileExistsError): - ensure_directories(mock_path, clobber=False) + ensure_directories("fakepath", clobber=False) mock_mkdir.assert_not_called() +@mock.patch("duct.__main__.Path", spec=pathlib.Path) @mock.patch("duct.__main__.os.makedirs") -@mock.patch("duct.__main__.glob.glob") def test_ensure_directories_with_conflicts_clobber( - mock_glob: mock.MagicMock, mock_mkdir: mock.MagicMock + mock_mkdir: mock.MagicMock, mock_path: mock.MagicMock ) -> None: - mock_path = "mock_dir/path" - mock_glob.return_value = ["possibly", "conflicting", "files"] - ensure_directories(mock_path, clobber=True) - mock_mkdir.assert_called_once_with("mock_dir", exist_ok=True) + mock_prefix = "fakepath/" + mock_path.return_value.exists.return_value = True + ensure_directories(mock_prefix, clobber=True) + mock_mkdir.assert_called_once_with(mock_prefix, exist_ok=True) @mock.patch("duct.__main__.os.makedirs") diff --git a/test/test_report.py b/test/test_report.py index e0d78480..f3eb3cac 100644 --- a/test/test_report.py +++ b/test/test_report.py @@ -1,7 +1,7 @@ from __future__ import annotations import subprocess from unittest.mock import MagicMock, patch -from duct.__main__ import ProcessStats, Report, Sample +from duct.__main__ import ProcessStats, Report, Sample, Suffix stat0 = ProcessStats( pcpu=0.0, pmem=0, rss=0, vsz=0, timestamp="2024-06-11T10:09:37-04:00" @@ -77,7 +77,7 @@ def test_write_pid_samples_set_first_run_only( report.write_pid_samples() assert report._resource_stats_log_path is not None # mypy says unreachable which is a lie. (_resource_stats_log_path set as side effect of write_pid_samples) - mock_open.assert_called_once_with(f"{temp_output_dir}usage.json", "a") # type: ignore + mock_open.assert_called_once_with(f"{temp_output_dir}{Suffix.usage}", "a") # type: ignore mock_clob_or_clear.reset_mock() report.write_pid_samples() mock_clob_or_clear.assert_not_called()