bec/bec_lib/tests/test_file_utils.py

""" 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,
    ServiceConfigParser,
)
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",
        info={"scan_number": 5, "file_suffix": "SampleA", "file_directory": "test_dir"},
        status="closed",
    )


def test_device_config_writer():
    """Device config writer fixture and ServiceConfigParser class"""
    with mock.patch("os.makedirs") as mock_make_dirs:
        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"


def test_log_writer():
    """Device config writer fixture and ServiceConfigParser class"""
    with mock.patch("os.makedirs") as mock_make_dirs:
        lw = LogWriter(service_config={"base_path": "/tmp"})
        assert mock_make_dirs.call_count == 1
        assert lw.directory == "/tmp/logs"


@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")
        suffix2 = suffix + f"_{scan_msg.info['file_suffix']}"
        expected = os.path.join(
            file_writer._base_path,
            "data",
            scan_msg.info["file_directory"],
            f"S{scannr:0{file_writer._leading_zeros}d}_{suffix2}.h5",
        )
        assert return_value == expected

    # case 3
    scan_msg.info.pop("file_suffix")
    scan_msg.info.pop("file_directory")
    with mock.patch.object(file_writer, "get_scan_msg", return_value=scan_msg):
        scannr = scan_msg.info.get("scan_number")
        scan_dir = f"S0000-0999/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