bec/bec_lib/tests/test_bec_service.py

import contextlib
import os
from pathlib import Path
from unittest import mock

import pytest

import bec_lib
from bec_lib import messages
from bec_lib.bec_service import BECService
from bec_lib.endpoints import MessageEndpoints
from bec_lib.logger import bec_logger
from bec_lib.messages import BECStatus
from bec_lib.service_config import ServiceConfig

# pylint: disable=no-member
# pylint: disable=missing-function-docstring
# pylint: disable=redefined-outer-name
# pylint: disable=protected-access

dir_path = os.path.dirname(bec_lib.__file__)


@contextlib.contextmanager
def bec_service(config, connector_cls=None, **kwargs):
    if connector_cls is None:
        connector_cls = mock.MagicMock()
    service = BECService(config=config, connector_cls=connector_cls, **kwargs)
    try:
        yield service
    finally:
        service.shutdown()
        bec_logger.logger.remove()
        bec_logger._reset_singleton()


def test_bec_service_init_with_service_config():
    config = ServiceConfig(redis={"host": "localhost", "port": 6379})
    with bec_service(config) as service:
        assert service._service_config == config
        assert service.bootstrap_server == "localhost:6379"
        assert service._unique_service is False


def test_bec_service_init_raises_for_invalid_config():
    with pytest.raises(TypeError):
        with bec_service(mock.MagicMock()):
            ...


def test_bec_service_init_with_service_config_path():
    with bec_service(config=f"{dir_path}/tests/test_service_config.yaml") as service:
        assert isinstance(service._service_config, ServiceConfig)
        assert service.bootstrap_server == "localhost:6379"
        assert service._unique_service is False


def test_init_runs_service_check():
    with mock.patch.object(
        BECService, "_update_existing_services", return_value=False
    ) as mock_update_existing_services:
        with bec_service(f"{dir_path}/tests/test_service_config.yaml", unique_service=True):
            mock_update_existing_services.assert_called_once()


def test_run_service_check_raises_for_existing_service():
    with mock.patch.object(
        BECService, "_update_existing_services", return_value=False
    ) as mock_update_existing_services:
        with bec_service(
            f"{dir_path}/tests/test_service_config.yaml", unique_service=True
        ) as service:
            service._services_info = {"BECService": mock.MagicMock()}
            with pytest.raises(RuntimeError):
                service._run_service_check(timeout_time=0, elapsed_time=10)


def test_run_service_check_repeats():
    with mock.patch.object(
        BECService, "_update_existing_services", return_value=False
    ) as mock_update_existing_services:
        with bec_service(
            f"{dir_path}/tests/test_service_config.yaml", unique_service=True
        ) as service:
            service._services_info = {"BECService": mock.MagicMock()}
            assert service._run_service_check(timeout_time=0.5, elapsed_time=0) is True


def test_bec_service_service_status():
    with mock.patch.object(
        BECService, "_update_existing_services", return_value=False
    ) as mock_update_existing_services:
        with bec_service(
            f"{dir_path}/tests/test_service_config.yaml", unique_service=True
        ) as service:
            mock_update_existing_services.reset_mock()
            status = service.service_status
            mock_update_existing_services.assert_called_once()


def test_bec_service_update_existing_services():
    service_keys = [
        MessageEndpoints.service_status("service1").endpoint.encode(),
        MessageEndpoints.service_status("service2").endpoint.encode(),
    ]
    service_msgs = [
        messages.StatusMessage(name="service1", status=BECStatus.RUNNING, info={}, metadata={}),
        messages.StatusMessage(name="service2", status=BECStatus.IDLE, info={}, metadata={}),
    ]
    connector_cls = mock.MagicMock()
    connector_cls().keys.return_value = service_keys
    connector_cls().get.side_effect = [msg for msg in service_msgs]
    with bec_service(
        f"{os.path.dirname(bec_lib.__file__)}/tests/test_service_config.yaml",
        connector_cls=connector_cls,
        unique_service=True,
    ) as service:
        assert service._services_info == {"service1": service_msgs[0], "service2": service_msgs[1]}


def test_bec_service_update_existing_services_ignores_wrong_msgs():
    service_keys = [
        MessageEndpoints.service_status("service1").endpoint.encode(),
        MessageEndpoints.service_status("service2").endpoint.encode(),
    ]
    service_msgs = [
        messages.StatusMessage(name="service1", status=BECStatus.RUNNING, info={}, metadata={}),
        None,
    ]
    connector_cls = mock.MagicMock()
    connector_cls().keys.return_value = service_keys
    connector_cls().get.side_effect = [service_msgs[0], None]
    with bec_service(
        f"{os.path.dirname(bec_lib.__file__)}/tests/test_service_config.yaml",
        connector_cls=connector_cls,
        unique_service=True,
    ) as service:
        assert service._services_info == {"service1": service_msgs[0]}


def test_bec_service_default_config():
    with bec_service(
        f"{os.path.dirname(bec_lib.__file__)}/tests/test_service_config.yaml", unique_service=True
    ) as service:
        assert service._service_config.service_config["file_writer"]["base_path"] == "./"

    config = ServiceConfig(redis={"host": "localhost", "port": 6379})
    with bec_service(config=config, unique_service=True) as service:
        bec_lib_path = str(Path(bec_lib.service_config.__file__).resolve().parent.parent.parent)
        if "nox" not in bec_lib_path:
            assert (
                os.path.abspath(service._service_config.service_config["file_writer"]["base_path"])
                == bec_lib_path
            )