85a371a38b
SIENTIAPDE-1325: Refactor Ingestor and Manager Classes for Enhanced Asynchronous Operations - Introduced asynchronous methods across Ingestor, IngestorManager, DataManager, and OpcManager to improve performance and responsiveness. - Integrated MetricsController into various classes for better observability and monitoring. - Updated Redis and MongoDB interactions to support asynchronous operations, enhancing data handling efficiency. - Removed deprecated Redis metrics and streamlined resource management logic. - Adjusted unit tests to accommodate the new asynchronous behavior and ensure proper mocking of async methods.
285 lines
9.1 KiB
Python
285 lines
9.1 KiB
Python
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
|
|
|
|
# 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'],
|
|
)
|