sientia-dataops-opc-ingestor/tests/unit/managers/test_data_manager.py

from unittest.mock import ANY, MagicMock, patch
from pytest import fixture
from kafka.errors import NoBrokersAvailable
from sientia_do.notifications.models import NotificationLevel
from ingestor.managers.data_manager import DataManager


@fixture
@patch("ingestor.managers.data_manager.KafkaProducer")
def data_manager(kafka):
    return DataManager(
        kafka_servers="localhost:9092",
        logger=MagicMock(),
        notification_handler=MagicMock()
    )


@patch("ingestor.managers.data_manager.KafkaProducer")
def test___init___success(kafka):
    logger_mock = MagicMock()

    data_manager = DataManager(
        kafka_servers="localhost:9092",
        logger=logger_mock,
        notification_handler=MagicMock()
    )

    kafka.assert_called_once_with(
        bootstrap_servers="localhost:9092",
        value_serializer=ANY,
        key_serializer=ANY
    )
    assert data_manager.kafka_producer is not None
    logger_mock.info.assert_any_call(
        "Trying (0) to initializing DataManager with Kafka servers: localhost:9092"
    )
    logger_mock.info.assert_any_call(
        "DataManager initialized with Kafka servers: localhost:9092"
    )
    logger_mock.error.assert_not_called()
    assert logger_mock.info.call_count == 2


@patch("ingestor.managers.data_manager.KafkaProducer")
def test___init___second_attempt(kafka):
    kafka.side_effect = [NoBrokersAvailable, MagicMock()]
    logger_mock = MagicMock()

    data_manager = DataManager(
        kafka_servers="localhost:9092",
        logger=logger_mock,
        notification_handler=MagicMock()
    )

    kafka.assert_any_call(
        bootstrap_servers="localhost:9092",
        value_serializer=ANY,
        key_serializer=ANY
    )
    assert kafka.call_count == 2
    assert data_manager.kafka_producer is not None
    logger_mock.info.assert_any_call(
        "Trying (0) to initializing DataManager with Kafka servers: localhost:9092"
    )
    logger_mock.info.assert_any_call(
        "Trying (1) to initializing DataManager with Kafka servers: localhost:9092"
    )
    logger_mock.info.assert_any_call(
        "DataManager initialized with Kafka servers: localhost:9092"
    )
    logger_mock.error.assert_called_once_with(
        "Kafka servers localhost:9092 are not available. Retrying..."
    )
    assert logger_mock.info.call_count == 3


@patch("ingestor.managers.data_manager.KafkaProducer")
def test___init___failure_max_attempts(kafka):
    kafka.side_effect = NoBrokersAvailable
    logger_mock = MagicMock()

    try:
        DataManager(
            kafka_servers="localhost:9092",
            logger=logger_mock,
            notification_handler=MagicMock()
        )
    except NoBrokersAvailable as e:
        assert str(
            e) == "NoBrokersAvailable: Failed to connect to Kafka servers localhost:9092 after 3 attempts."

        assert kafka.call_count == 3
        logger_mock.info.assert_any_call(
            "Trying (0) to initializing DataManager with Kafka servers: localhost:9092"
        )
        logger_mock.info.assert_any_call(
            "Trying (1) to initializing DataManager with Kafka servers: localhost:9092"
        )
        logger_mock.info.assert_any_call(
            "Trying (2) to initializing DataManager with Kafka servers: localhost:9092"
        )
        logger_mock.error.assert_called_with(
            "Failed to connect to Kafka servers localhost:9092 after 3 attempts."
        )
        assert logger_mock.info.call_count == 3

    else:
        assert False, "Expected NoBrokersAvailable exception was not raised."


def test___del___has_producer(data_manager):
    flush_mock = MagicMock()
    close_mock = MagicMock()

    data_manager.kafka_producer.flush = flush_mock
    data_manager.kafka_producer.close = close_mock

    data_manager.__del__()
    flush_mock.assert_called_once()
    close_mock.assert_called_once()


@patch("ingestor.managers.data_manager.print")
def test___del___no_producer(print, data_manager):
    data_manager.kafka_producer = None

    # Call the __del__ method
    data_manager.__del__()

    # Check if the print statement was called
    print.assert_any_call(
        "Kafka producer is already closed or not initialized."
    )


def test_delivery_report(data_manager):
    msg = MagicMock()
    msg.topic = "test_topic"
    msg.partition = 0
    msg.offset = 1

    data_manager.delivery_report(msg)

    data_manager.logger.debug.assert_called_once_with(
        f"Record successfully produced to {msg.topic} [{msg.partition}] at offset {msg.offset}"
    )


def test_delivery_error(data_manager):
    err = "Test error"
    data_manager.delivery_error(err)

    data_manager.logger.error.assert_called_once_with(
        f"Delivery failed for record : {err}"
    )


def test_publish(data_manager):
    topic = "test_topic"
    data = {"key": "value"}

    # Mock the send method of the Kafka producer
    send_mock = MagicMock()
    data_manager.kafka_producer.send = send_mock

    # Call the publish method
    data_manager.publish(topic, data)

    # Check if the send method was called with the correct arguments
    send_mock.assert_called_once_with(
        topic=topic, value=data
    )

    send_mock.return_value.add_callback.assert_called_once()

    data_manager.kafka_producer.flush.assert_called_once()


@patch("ingestor.managers.data_manager.traceback")
def test_publish_error(traceback, data_manager):
    topic = "test_topic"
    data = {"key": "value"}

    # Mock the send method of the Kafka producer to raise an exception
    send_mock = MagicMock(side_effect=Exception("Test error"))
    data_manager.kafka_producer.send = send_mock

    # Call the publish method
    data_manager.publish(topic, data)

    # Check if the send method was called with the correct arguments
    send_mock.assert_called_once_with(
        topic=topic, value=data
    )

    # Check if the error was logged
    data_manager.notification_handler.build_and_send_notification.assert_called_once_with(
        notification_id=f"KAFKA_PRODUCER_ERROR_{topic}",
        message=f"Error publishing message to topic {topic}: Test error",
        block="kafka_producer",
        level=NotificationLevel.ERROR,
        attachment_content=traceback.format_exc.return_value
    )