1a96ada0e0
Refactor test assertions for improved readability and consistency - Updated assertions in test cases for DataManager and OpcManager to enhance readability by formatting long lines. - Ensured that the expected behavior of metrics logging and message publishing is clearly defined in the tests.
281 lines
8.5 KiB
Python
281 lines
8.5 KiB
Python
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")
|
|
@patch("ingestor.managers.data_manager.MongoClient")
|
|
def data_manager(mongo, kafka):
|
|
return DataManager(
|
|
kafka_servers="localhost:9092",
|
|
mongo_connection_string="mongodb://localhost:27017",
|
|
mongo_database="sientia",
|
|
export_to_kafka=True,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock()
|
|
)
|
|
|
|
|
|
@patch("ingestor.managers.data_manager.KafkaProducer")
|
|
@patch("ingestor.managers.data_manager.MongoClient")
|
|
def test___init___success(mongo, kafka):
|
|
logger_mock = MagicMock()
|
|
|
|
data_manager = DataManager(
|
|
kafka_servers="localhost:9092",
|
|
mongo_connection_string="mongodb://localhost:27017",
|
|
mongo_database="sientia",
|
|
export_to_kafka=True,
|
|
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 == 4
|
|
|
|
|
|
@patch("ingestor.managers.data_manager.KafkaProducer")
|
|
@patch("ingestor.managers.data_manager.MongoClient")
|
|
def test___init___second_attempt(mongo, kafka):
|
|
kafka.side_effect = [NoBrokersAvailable, MagicMock()]
|
|
logger_mock = MagicMock()
|
|
|
|
data_manager = DataManager(
|
|
kafka_servers="localhost:9092",
|
|
mongo_connection_string="mongodb://localhost:27017",
|
|
mongo_database="sientia",
|
|
export_to_kafka=True,
|
|
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 == 5
|
|
|
|
|
|
@patch("ingestor.managers.data_manager.KafkaProducer")
|
|
@patch("ingestor.managers.data_manager.MongoClient")
|
|
def test___init___failure_max_attempts(mongo, kafka):
|
|
kafka.side_effect = NoBrokersAvailable
|
|
logger_mock = MagicMock()
|
|
|
|
try:
|
|
DataManager(
|
|
kafka_servers="localhost:9092",
|
|
mongo_connection_string="mongodb://localhost:27017",
|
|
mongo_database="sientia",
|
|
export_to_kafka=True,
|
|
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_shutdown_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.shutdown()
|
|
flush_mock.assert_called_once()
|
|
close_mock.assert_called_once()
|
|
|
|
|
|
def test_shutdown_no_producer(data_manager):
|
|
data_manager.kafka_producer = None
|
|
|
|
data_manager.shutdown()
|
|
|
|
data_manager.logger.warning.assert_any_call(
|
|
"Kafka producer is already closed or not initialized."
|
|
)
|
|
|
|
|
|
def test_shutdown_no_mongo_client(data_manager):
|
|
data_manager.mongo_client = None
|
|
|
|
data_manager.shutdown()
|
|
|
|
data_manager.logger.warning.assert_any_call(
|
|
"MongoDB client is already closed or not initialized."
|
|
)
|
|
|
|
|
|
def test_shutdown_exception(data_manager):
|
|
data_manager.kafka_producer.flush = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
data_manager.kafka_producer.close = MagicMock()
|
|
|
|
data_manager.shutdown()
|
|
data_manager.logger.error.assert_called_once_with(
|
|
"Error closing Kafka producer: Test error"
|
|
)
|
|
|
|
|
|
def test_shutdown_exception_mongo(data_manager):
|
|
data_manager.mongo_client.close = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
|
|
data_manager.shutdown()
|
|
|
|
data_manager.logger.error.assert_called_once_with(
|
|
"Error closing MongoDB client: Test error"
|
|
)
|
|
|
|
|
|
def test___del__(data_manager):
|
|
data_manager.shutdown = MagicMock()
|
|
data_manager.__del__()
|
|
data_manager.shutdown.assert_called_once()
|
|
|
|
|
|
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()
|
|
|
|
|
|
def test_publish_no_kafka(data_manager):
|
|
data_manager.export_to_kafka = False
|
|
topic = "test_topic"
|
|
data = {"key": "value"}
|
|
|
|
data_manager.publish(topic, data)
|
|
|
|
data_manager.kafka_producer.send.assert_not_called()
|
|
|
|
|
|
@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
|
|
)
|
|
|
|
|
|
def test_publish_error_mongo(data_manager):
|
|
data_manager.export_to_kafka = False
|
|
data_manager.mongo_db.__getitem__.return_value.insert_one = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
|
|
data_manager.publish("test_topic", {"key": "value"})
|
|
|
|
data_manager.notification_handler.build_and_send_notification.assert_called_once_with(
|
|
notification_id="MONGO_PRODUCER_ERROR_test_topic",
|
|
message="Error inserting message to MongoDB: Test error",
|
|
block="mongo_producer",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|