"""Module to test file_utils.py""" # pylint: skip-file import os from unittest import mock import pytest from bec_lib.bec_errors import ServiceConfigError from bec_lib.endpoints import MessageEndpoints from bec_lib.file_utils import ( DeviceConfigWriter, FileWriter, FileWriterError, LogWriter, compile_file_components, get_full_path, ) from bec_lib.messages import ScanStatusMessage from bec_lib.tests.utils import ConnectorMock @pytest.fixture(scope="function") def mock_connector(): """Mock connector fixture""" connector = ConnectorMock("") yield connector @pytest.fixture(scope="function") def file_writer(mock_connector): """File writer fixture""" with mock.patch("os.makedirs") as mock_make_dirs: yield FileWriter(service_config={"base_path": "/tmp"}, connector=mock_connector) @pytest.fixture(scope="function") def scan_msg(): """Scan message fixture""" yield ScanStatusMessage( scan_id="1111", scan_parameters={"system_config": {"file_directory": None, "file_suffix": None}}, info={"scan_number": 5, "file_components": ("S00000-00999/S00005/S00005", "h5")}, status="closed", ) def test_device_config_writer(): """Device config writer fixture and ServiceConfigParser class""" with mock.patch("os.makedirs") as mock_make_dirs: with mock.patch("os.chmod") as mock_chmod: dcw = DeviceConfigWriter(service_config={"base_path": "/tmp"}) assert mock_make_dirs.call_count == 1 assert dcw.directory == "/tmp/device_configs" assert dcw.get_recovery_directory() == "/tmp/device_configs/recovery_configs" mock_chmod.assert_called_once_with("/tmp/device_configs", int("0o771", 8)) def test_log_writer(): """Device config writer fixture and ServiceConfigParser class""" with mock.patch("os.makedirs") as mock_make_dirs: with mock.patch("os.chmod") as mock_chmod: lw = LogWriter(service_config={"base_path": "/tmp"}) assert mock_make_dirs.call_count == 1 assert lw.directory == "/tmp/logs" mock_chmod.assert_called_once_with("/tmp/logs", int("0o771", 8)) @pytest.mark.parametrize( "service_config, connector, raises", [ ({"base_path": "/tmp"}, ConnectorMock(""), False), ({"base_path": "/tmp"}, None, False), ({"base_path": None}, ConnectorMock(""), True), ], ) def test_file_writer_init(service_config, connector, raises): """Test file writer init""" if raises: with pytest.raises(ServiceConfigError): with mock.patch("os.makedirs") as mock_make_dirs: fw = FileWriter(service_config=service_config, connector=connector) else: with mock.patch("os.makedirs") as mock_make_dirs: fw = FileWriter(service_config=service_config, connector=ConnectorMock("")) if service_config.get("base_path"): assert fw._base_path == service_config.get("base_path") assert fw._configured is True else: assert fw._base_path == "." assert fw._configured is False def test_get_scan_msg(file_writer, scan_msg): """Test get_scan_msg method""" file_writer.connector._get_buffer = {MessageEndpoints.scan_status().endpoint: scan_msg} rtr = file_writer.get_scan_msg() assert rtr == scan_msg @pytest.mark.parametrize( "scan_number,bundle,lead,user_suffix,ref_path", [ (10, 1000, None, None, "S0000-0999/S0010"), (2001, 1000, None, "", "S2000-2999/S2001"), (20, 50, 4, "sampleA", "S0000-0049/S0020_sampleA"), (20, 50, 5, "sampleB", "S00000-00049/S00020_sampleB"), (1200, 1000, 5, "", "S01000-01999/S01200"), ], ) def test_file_writer_get_scan_dir(scan_number, bundle, lead, user_suffix, ref_path, file_writer): dir_path = file_writer.get_scan_directory( scan_bundle=bundle, scan_number=scan_number, leading_zeros=lead, user_suffix=user_suffix ) assert dir_path == ref_path def test_compile_full_filename_not_configured(file_writer, scan_msg): file_writer.get_scan_msg = mock.MagicMock(return_value=scan_msg) file_writer._configured = False suffix = "test" path = os.path.join(file_writer._base_path, "data", f"S00000_default_{suffix}.h5") fp = file_writer.compile_full_filename(suffix=suffix) assert fp == path def test_compile_full_filename(file_writer, scan_msg): suffix = "test" file_type = ".csv" # case 1 with mock.patch.object(file_writer, "get_scan_msg", return_value=None): return_value = file_writer.compile_full_filename(suffix=suffix) assert return_value == "" # case 2 with mock.patch.object(file_writer, "get_scan_msg", return_value=scan_msg): return_value = file_writer.compile_full_filename(suffix=suffix) scannr = scan_msg.info.get("scan_number") expected = os.path.join( file_writer._base_path, "data", f"{scan_msg.info['file_components'][0]}_{suffix}.{scan_msg.info['file_components'][1]}", ) assert return_value == expected # case 3 scan_msg.scan_parameters.pop("system_config") with mock.patch.object(file_writer, "get_scan_msg", return_value=scan_msg): scannr = scan_msg.info.get("scan_number") scan_dir = f"S00000-00999/S{scannr:0{file_writer._leading_zeros}d}" expected = os.path.join( file_writer._base_path, "data", scan_dir, f"S{scannr:0{file_writer._leading_zeros}d}_{suffix}.h5", ) with mock.patch.object(file_writer, "get_scan_directory", return_value=scan_dir): return_value = file_writer.compile_full_filename(suffix=suffix) assert return_value == expected def test_compile_file_components(): """Test compile_file_components method. Case 1: test default values Case 2: test with file_directory and file_suffix (file suffix will not play a role) Case 3: test with file suffix, no file directory """ basepath = "/tmp" scan_nr = 10 file_path, extension = compile_file_components(basepath, scan_nr) assert extension == "h5" assert file_path == "/tmp/data/S00000-00999/S00010/S00010" scan_bundle = 1000 leading_zeros = 5 file_directory = "test_dir" user_suffix = "SampleA" file_path, extension = compile_file_components( base_path=basepath, scan_nr=scan_nr, scan_bundle=scan_bundle, leading_zeros=leading_zeros, file_directory=file_directory, user_suffix=user_suffix, ) assert extension == "h5" assert file_path == "/tmp/data/test_dir/S00010" file_path, extension = compile_file_components( base_path=basepath, scan_nr=scan_nr, scan_bundle=scan_bundle, leading_zeros=leading_zeros, user_suffix=user_suffix, ) assert extension == "h5" assert file_path == "/tmp/data/S00000-00999/S00010_SampleA/S00010" @pytest.mark.parametrize( "scan_info", [ ( ScanStatusMessage( scan_id="1111", scan_parameters={"system_config": {"file_directory": None, "file_suffix": None}}, info={"scan_number": 5, "file_components": ("S00000-00999/S00005/S00005", "h5")}, status="closed", ) ), ( ScanStatusMessage( scan_id="1111", scan_parameters={ "system_config": {"file_directory": "/my_dir", "file_suffix": None} }, info={"scan_number": 5, "file_components": ("/my_dir/S00005", "h5")}, status="closed", ) ), ( ScanStatusMessage( scan_id="1111", scan_parameters={ "system_config": {"file_directory": None, "file_suffix": "sampleA"} }, info={ "scan_number": 5, "file_components": ("S00000-00999/S00005_sampleA/S00005", "h5"), }, status="closed", ) ), ], ) @pytest.mark.parametrize( "name, expect_error", [ ("valid_name", False), # Valid case ("valid-name", False), # Valid case with '-' ("valid_name123", False), # Valid case with '_' ("", True), # Empty string should raise an error ("invalid name", True), # Contains space ("name$", True), # Contains non-alphanumeric character ("名前", True), # Non-ASCII characters ("1234", False), # Numeric values are allowed ("_hidden", False), # Leading underscore is allowed ("-dashprefix", False), # Leading dash is allowed ], ) def test_get_full_path(scan_info, name, expect_error): file_components = scan_info.info.get("file_components") suffix = scan_info.scan_parameters.get("system_config").get("file_suffix") if expect_error: with pytest.raises(FileWriterError): get_full_path(scan_info, name, create_dir=False) else: if suffix is None: full_path = f"{file_components[0]}_{name}.{file_components[1]}" else: full_path = f"{file_components[0]}_{name}_{suffix}.{file_components[1]}" ret = get_full_path(scan_info, name, create_dir=False) assert ret == full_path