Code import - branch main
This commit is contained in:
283
tests/unit/managers/test_data_manager.py
Normal file
283
tests/unit/managers/test_data_manager.py
Normal file
@@ -0,0 +1,283 @@
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
|
||||
from kafka.errors import NoBrokersAvailable
|
||||
from pytest import fixture, mark
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
from ingestor.managers.data_manager import DataManager
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
'pod_id': 'localhost',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('ingestor.managers.data_manager.KafkaProducer')
|
||||
@patch('ingestor.managers.data_manager.MongoDBRepository')
|
||||
def data_manager(mongodb_repository, kafka):
|
||||
data_manager = DataManager(
|
||||
kafka_servers='localhost:9092',
|
||||
mongo_connection_string='mongodb://localhost:27017',
|
||||
mongo_database='sientia',
|
||||
export_to_kafka=True,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metadata=metadata['metadata'],
|
||||
metrics_controller=MagicMock(),
|
||||
)
|
||||
|
||||
data_manager.send_notification = MagicMock()
|
||||
data_manager.send_notification_async = AsyncMock()
|
||||
data_manager.emit_metric = AsyncMock()
|
||||
return data_manager
|
||||
|
||||
|
||||
@patch('ingestor.managers.data_manager.KafkaProducer')
|
||||
@patch('ingestor.managers.data_manager.MongoDBRepository')
|
||||
def test___init___success(mongodb_repository, kafka):
|
||||
logger_mock = MagicMock()
|
||||
|
||||
data_manager = DataManager(
|
||||
metadata=metadata['metadata'],
|
||||
kafka_servers='localhost:9092',
|
||||
mongo_connection_string='mongodb://localhost:27017',
|
||||
mongo_database='sientia',
|
||||
export_to_kafka=True,
|
||||
logger=logger_mock,
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=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.MongoDBRepository')
|
||||
def test___init___second_attempt(mongodb_repository, 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(),
|
||||
metadata=metadata['metadata'],
|
||||
metrics_controller=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.MongoDBRepository')
|
||||
def test___init___failure_max_attempts(mongodb_repository, 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(),
|
||||
metadata=metadata['metadata'],
|
||||
metrics_controller=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:
|
||||
raise AssertionError('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_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_repository.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}')
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async 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
|
||||
await 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')
|
||||
@mark.asyncio
|
||||
async def test_publish_error(traceback, data_manager):
|
||||
topic = 'test_topic'
|
||||
data = {'key': 'value', 'name': 'test_tag'}
|
||||
|
||||
# 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
|
||||
|
||||
data_manager.mongo_repository = AsyncMock()
|
||||
|
||||
# Call the publish method
|
||||
await 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.send_notification_async.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,
|
||||
metadata=metadata['metadata'],
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_publish_error_mongo(data_manager):
|
||||
data_manager.export_to_kafka = False
|
||||
data_manager.mongo_repository.insert = AsyncMock(side_effect=Exception('Test error'))
|
||||
|
||||
await data_manager.publish('test_topic', {'key': 'value'})
|
||||
|
||||
data_manager.send_notification_async.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,
|
||||
metadata=metadata['metadata'],
|
||||
)
|
||||
Reference in New Issue
Block a user