import contextlib import os import time 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={}), ] service_metric_msgs = [ messages.ServiceMetricMessage(name="service1", metrics={}), messages.ServiceMetricMessage(name="service2", metrics={}), ] connector_cls = mock.MagicMock() connector_cls().keys.return_value = service_keys msgs = service_msgs + service_metric_msgs connector_cls().get.side_effect = [msg for msg in 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]} assert service._services_metric == { "service1": service_metric_msgs[0], "service2": service_metric_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, ] service_metric_msgs = [None, messages.ServiceMetricMessage(name="service2", metrics={})] msgs = service_msgs + service_metric_msgs connector_cls = mock.MagicMock() connector_cls().keys.return_value = service_keys connector_cls().get.side_effect = [msg for msg in 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]} 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 ) def test_bec_service_show_global_vars(capsys): config = ServiceConfig(redis={"host": "localhost", "port": 6379}) with bec_service(config=config, unique_service=True) as service: ep = MessageEndpoints.global_vars("test").endpoint.encode() with mock.patch.object(service.connector, "keys", return_value=[ep]): with mock.patch.object(service, "get_global_var", return_value="test_value"): service.show_global_vars() captured = capsys.readouterr() assert "test" in captured.out assert "test_value" in captured.out def test_bec_service_globals(connected_connector): config = ServiceConfig(redis={"host": "localhost", "port": 1}) with bec_service(config=config, unique_service=True) as service: service.connector = connected_connector service.set_global_var("test", "test_value") assert service.get_global_var("test") == "test_value" service.delete_global_var("test") assert service.get_global_var("test") is None def test_bec_service_metrics(connected_connector): config = ServiceConfig(redis={"host": "localhost", "port": 1}) with mock.patch("bec_lib.bec_service.BECService._start_metrics_emitter") as mock_emitter: with bec_service(config=config, unique_service=True) as service: service._metrics_emitter_event = mock.MagicMock() service._metrics_emitter_event.is_set.side_effect = [False, True] service.connector = connected_connector assert service._services_metric == {} service._send_service_status() service._get_metrics() service._update_existing_services() assert service._services_metric != {}