diff --git a/pslab/cli.py b/pslab/cli.py index 5ba833b..e960f0b 100644 --- a/pslab/cli.py +++ b/pslab/cli.py @@ -21,6 +21,7 @@ import struct import sys import time +from contextlib import suppress from itertools import zip_longest from typing import List, Tuple @@ -146,24 +147,26 @@ def collect(handler: SerialHandler, args: argparse.Namespace): output = instrument(handler, args.channels, args.duration) - if args.file_path is not None: - file = open(args.file_path, "w") - else: - file = sys.stdout - - if not args.json: - csv_file = csv.writer(file) - csv_file.writerow(output[0]) - for row in zip_longest(*output[1]): - csv_file.writerow(row) + file = open(args.file_path, "w") if args.file_path is not None else sys.stdout + try: + if not args.json: + csv_file = csv.writer(file) + csv_file.writerow(output[0]) + for row in zip_longest(*output[1]): + csv_file.writerow(row) + else: + output_dict = dict() + for key, val in zip_longest(*output): + output_dict[key] = val.tolist() + json.dump(output_dict, file) + except BaseException: + if args.file_path is not None: + with suppress(OSError): + file.close() + raise else: - output_dict = dict() - for key, val in zip_longest(*output): - output_dict[key] = val.tolist() - json.dump(output_dict, file) - - if args.file_path is not None: - file.close() + if args.file_path is not None: + file.close() def wave(handler: SerialHandler, args: argparse.Namespace): diff --git a/tests/test_cli_output_cleanup.py b/tests/test_cli_output_cleanup.py new file mode 100644 index 0000000..21a772e --- /dev/null +++ b/tests/test_cli_output_cleanup.py @@ -0,0 +1,134 @@ +"""CLI output resources are released without closing caller-owned streams.""" + +import argparse +import builtins +import io +import json + +import numpy as np +import pytest + +from pslab import cli + + +class FailingOutput: + """Delegate file ownership while simulating a failed output write.""" + + def __init__(self, file): + self.file = file + self.error = OSError("output storage unavailable") + + def write(self, text): + self.file.write(text[:1]) + raise self.error + + def close(self): + self.file.close() + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + +class FailingCloseOutput(FailingOutput): + """Expose a close failure after releasing the underlying file.""" + + def __init__(self, file): + super().__init__(file) + self.close_error = OSError("output close failed") + + def close(self): + super().close() + raise self.close_error + + +@pytest.fixture +def output_args(monkeypatch): + monkeypatch.setitem( + cli.INSTRUMENTS, + "test-instrument", + lambda *args: (["CH1"], [np.array([1.0, 2.0])]), + ) + return argparse.Namespace( + instrument="test-instrument", channels=1, duration=1, file_path=None, json=False + ) + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_failed_file_output_is_closed(tmp_path, monkeypatch, output_args, as_json): + """Release the real file even when the write exception retains its frame.""" + destination = tmp_path / "capture" + file = builtins.open(destination, "w") + output = FailingOutput(file) + monkeypatch.setattr(cli, "open", lambda *args: output, raising=False) + output_args.file_path = str(destination) + output_args.json = as_json + try: + with pytest.raises(OSError) as raised: + cli.collect(None, output_args) + assert raised.value is output.error + assert file.closed + finally: + file.close() + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_failed_stdout_is_not_closed(monkeypatch, output_args, as_json): + """Preserve caller ownership after either formatter encounters a write error.""" + file = io.StringIO() + output = FailingOutput(file) + monkeypatch.setattr(cli.sys, "stdout", output) + output_args.json = as_json + with pytest.raises(OSError) as raised: + cli.collect(None, output_args) + assert raised.value is output.error + assert not file.closed + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_successful_file_output_is_closed(tmp_path, monkeypatch, output_args, as_json): + """Close successful output while retaining CSV and JSON content.""" + destination = tmp_path / "capture" + file = builtins.open(destination, "w") + monkeypatch.setattr(cli, "open", lambda *args: file, raising=False) + output_args.file_path = str(destination) + output_args.json = as_json + cli.collect(None, output_args) + assert file.closed + if as_json: + assert json.loads(destination.read_text()) == {"CH1": [1.0, 2.0]} + else: + assert destination.read_text().splitlines() == ["CH1", "1.0", "2.0"] + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_close_failure_does_not_replace_write_error( + tmp_path, monkeypatch, output_args, as_json +): + file = builtins.open(tmp_path / "capture", "w") + output = FailingCloseOutput(file) + monkeypatch.setattr(cli, "open", lambda *args: output, raising=False) + output_args.file_path = str(tmp_path / "capture") + output_args.json = as_json + with pytest.raises(OSError) as raised: + cli.collect(None, output_args) + assert raised.value is output.error + assert file.closed + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_close_failure_after_successful_write_is_visible( + tmp_path, monkeypatch, output_args, as_json +): + file = builtins.open(tmp_path / "capture", "w") + output = FailingCloseOutput(file) + monkeypatch.setattr(output, "write", file.write) + monkeypatch.setattr(cli, "open", lambda *args: output, raising=False) + output_args.file_path = str(tmp_path / "capture") + output_args.json = as_json + with pytest.raises(OSError) as raised: + cli.collect(None, output_args) + assert raised.value is output.close_error + assert file.closed