From cf74ea135115a542fdc42a94d990996f45a77a1d Mon Sep 17 00:00:00 2001 From: Julian Stirling Date: Sun, 3 Aug 2025 21:29:22 +0100 Subject: [PATCH] More complete testing of code that runs stitching subprocess. --- .../stitching.py | 27 ++++- .../utilities.py | 1 + tests/mock_stitching/mock-stitch.py | 19 +++ tests/test_stitching.py | 110 +++++++++++++++++- 4 files changed, 151 insertions(+), 6 deletions(-) create mode 100644 tests/mock_stitching/mock-stitch.py diff --git a/src/openflexure_microscope_server/stitching.py b/src/openflexure_microscope_server/stitching.py index ec4fe5b6..47c61cca 100644 --- a/src/openflexure_microscope_server/stitching.py +++ b/src/openflexure_microscope_server/stitching.py @@ -10,6 +10,7 @@ from typing import Optional, Any import threading import subprocess import os +import shlex import time import labthings_fastapi as lt @@ -27,6 +28,21 @@ class StitcherValidationError(RuntimeError): """The stitcher received values that it deems unsafe to create a command from.""" +def validate_command(cmd: list[str]): + """Validate that the command only characters that are allowed in a path. + + The values in the commands should be numbers, commandline flags, paths, and + executables. All of these should be allowed by ``make_path_safe``. + + :raises StitcherValidationError: if any element in the command is not safe. + """ + for element in cmd: + if element != make_path_safe(element): + raise StitcherValidationError( + "Invalid stiching command: Contains unsafe characters." + ) + + class BaseStitcher: """A base stitching class for all stitchers. Don't initialise this directly.""" @@ -62,8 +78,9 @@ class BaseStitcher: """The command to run with subprocess.Popen.""" # Revalidate for good measure, self.validate_path() + # The command, and the mode - initial_args = [STITCHING_CMD, "--stitching_mode", self._mode] + initial_args = shlex.split(STITCHING_CMD) + ["--stitching_mode", self._mode] # Use float() just to really ensure that anything input is a float setting_args = [ "--minimum_overlap", @@ -71,7 +88,9 @@ class BaseStitcher: "--resize", f"{float(self.correlation_resize)}", ] - return initial_args + self._extra_args + setting_args + [self.images_dir] + full_cmd = initial_args + self._extra_args + setting_args + [self.images_dir] + validate_command(full_cmd) + return full_cmd def validate_path(self): """Check path is safe before making a command to run with subprocess. @@ -82,7 +101,7 @@ class BaseStitcher: """ if self.images_dir != make_path_safe(self.images_dir): raise StitcherValidationError( - "Invalid directory path contains unsafe characters." + "Invalid directory path: Contains unsafe characters." ) def start(self) -> None: @@ -280,7 +299,7 @@ class FinalStitcher(BaseStitcher): if process.poll() == 0: self.logger.info("Stitching complete") else: - raise ChildProcessError(f"Subprocess {cmd[0]} exited with an error.") + raise ChildProcessError(f"Subprocess {STITCHING_CMD} exited with an error.") def _log_ongoing( self, diff --git a/src/openflexure_microscope_server/utilities.py b/src/openflexure_microscope_server/utilities.py index 1031aa23..07bc4fea 100644 --- a/src/openflexure_microscope_server/utilities.py +++ b/src/openflexure_microscope_server/utilities.py @@ -144,6 +144,7 @@ def make_path_safe(unsafe_path_string: str) -> str: slash/backslash/colon on Windows). :param unsafe_path_string: The original path string to sanitise. + :returns: A version of the input string safe to use as a file path. """ unsafe_character_pattern = ( diff --git a/tests/mock_stitching/mock-stitch.py b/tests/mock_stitching/mock-stitch.py new file mode 100644 index 00000000..b4c9d16a --- /dev/null +++ b/tests/mock_stitching/mock-stitch.py @@ -0,0 +1,19 @@ +"""CLI script used for testing subprocess calls that should go to openflexure-stitch.""" + +import sys +import time + + +def main(): + """Echo command-line arguments with a short delay between each.""" + input_arguments = sys.argv[1:] + for arg in input_arguments: + # This is used to check we catch errors correctly. + if arg == "ERROR": + raise RuntimeError("I was told to do this.") + print(arg, flush=True) + time.sleep(0.2) + + +if __name__ == "__main__": + main() diff --git a/tests/test_stitching.py b/tests/test_stitching.py index 92189bd7..a680f71f 100644 --- a/tests/test_stitching.py +++ b/tests/test_stitching.py @@ -7,9 +7,15 @@ generated, and the subprocess calling works as expected. import logging import os from copy import copy +import uuid +import re +import threading +from time import sleep import pytest +import labthings_fastapi as lt + from openflexure_microscope_server.stitching import ( BaseStitcher, PreviewStitcher, @@ -20,7 +26,9 @@ from openflexure_microscope_server.stitching import ( # A global logger pretending to be an Invocation Logger LOGGER = logging.getLogger("mock-invocation_logger") -FAKE_DIR = os.path.join("a", "dir", "that", "is", "fake") +FAKE_DIR: list[str] = os.path.join("a", "dir", "that", "is", "fake") +THIS_DIR: str = os.path.dirname(os.path.realpath(__file__)) +MOCK_STITCHER: str = os.path.join(THIS_DIR, "mock_stitching", "mock-stitch.py") def test_base_stitcher(): @@ -165,7 +173,7 @@ def _validation_error_tester(scan_path, **kwargs): def test_validation_error(): - """Test a number of way to try to inject mallicious arguments into the stitcher. + """Test a number of way to try to inject malicious arguments into the stitcher. The stitcher should throw a validation error each attempt. """ @@ -173,3 +181,101 @@ def test_validation_error(): _validation_error_tester(FAKE_DIR, overlap=".2", correlation_resize=".25;rm -rf /;") _validation_error_tester(FAKE_DIR, overlap=".2;rm -rf /;", correlation_resize=".25") _validation_error_tester(FAKE_DIR, scan_data_dict={"overlap": ".2;rm -rf /;"}) + + +def test_extra_arg_validation(): + """Test that malicious arguments in extra_args also throw validation error. + + Currently extra args do not come from user input. But this makes checks more + future-proof. + """ + stitcher = FinalStitcher(FAKE_DIR, logger=LOGGER) + stitcher._extra_args = ["&&rm -rf /&&"] + with pytest.raises(StitcherValidationError): + stitcher.command + + +def test_preview_stitching_command(caplog, mocker): + """Check the preview process runs in a background thread and doesn't log.""" + mock_cmd = "python -m mock_command.py" + + mocker.patch("openflexure_microscope_server.stitching.STITCHING_CMD", mock_cmd) + + with caplog.at_level(logging.INFO): + stitcher = PreviewStitcher(FAKE_DIR, overlap=0.1, correlation_resize=0.5) + stitcher.start() + # Should take a second or so to run so will still be running + assert stitcher.running + # Can't start another time, instead get a runtime error + with pytest.raises(RuntimeError): + stitcher.start() + # Wait for it to complete + stitcher.wait() + # It is now not running + assert not stitcher.running + assert len(caplog.records) == 0 + + +def test_final_stitching_command(caplog, mocker): + """Check the final stitch runs until completion, and print statements are logged.""" + mock_cmd = f"python {MOCK_STITCHER}" + + mocker.patch("openflexure_microscope_server.stitching.STITCHING_CMD", mock_cmd) + + with caplog.at_level(logging.INFO): + # Input values to prevent logging + stitcher = FinalStitcher( + FAKE_DIR, logger=LOGGER, overlap=0.1, correlation_resize=0.5 + ) + # For the final stitcher it will always complete before returning. + stitcher.start(cancel=lt.deps.CancelHook(id=uuid.uuid4())) + # The mock command logs the inputs so should be 1 less than + # FINAL_EXPECTED_COMMAND as the command itself is not logged. However, there + # is two extra logs, The date (before the program starts) and + # "Stitching complete" when it ends. + assert len(caplog.records) == len(FINAL_EXPECTED_COMMAND) + 1 + for i, record in enumerate(caplog.records): + msg = record.message + if i == 0: + assert re.match(r"^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}", msg) + elif i == len(FINAL_EXPECTED_COMMAND): + assert msg.strip() == "Stitching complete" + else: + assert msg.strip() == FINAL_EXPECTED_COMMAND[i] + + +def test_final_stitching_command_cancelled(caplog, mocker): + """Check that final stitch can be cancelled.""" + mock_cmd = f"python {MOCK_STITCHER}" + + mocker.patch("openflexure_microscope_server.stitching.STITCHING_CMD", mock_cmd) + + with caplog.at_level(logging.INFO): + stitcher = FinalStitcher(FAKE_DIR, logger=LOGGER) + cancel_hook = lt.deps.CancelHook(id=uuid.uuid4()) + + # Start stitching in a thread. + thread = threading.Thread(target=stitcher.start, kwargs={"cancel": cancel_hook}) + thread.start() + # Sleep long enough for at least 1 log. + sleep(0.5) + # use set() to cancel! + cancel_hook.set() + thread.join() + assert len(caplog.records) < len(FINAL_EXPECTED_COMMAND) + 1 + assert caplog.records[-1].message == "Stitching cancelled by user" + + +def test_final_stitching_command_error(caplog, mocker): + """Check that ChildProcessError is raised if the final stitch errors.""" + mock_cmd = f"python {MOCK_STITCHER}" + + mocker.patch("openflexure_microscope_server.stitching.STITCHING_CMD", mock_cmd) + + with caplog.at_level(logging.INFO): + stitcher = FinalStitcher(FAKE_DIR, logger=LOGGER) + # Send in the argument ERROR to mock-stitch and it will raise an error rather + # than echo. + stitcher._extra_args = ["ERROR"] + with pytest.raises(ChildProcessError): + stitcher.start(cancel=lt.deps.CancelHook(id=uuid.uuid4()))