Code import - branch main

This commit is contained in:
2026-06-28 03:02:59 +00:00
commit 24af3bc6e3
38 changed files with 6722 additions and 0 deletions

View File

View 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'],
)

View File

@@ -0,0 +1,654 @@
from unittest.mock import AsyncMock, MagicMock, call, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from ingestor.managers.ingestor_manager import IngestorManager
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@fixture
@patch('ingestor.managers.ingestor_manager.DataManager')
@patch('ingestor.managers.ingestor_manager.ResourceManager')
def ingestor_manager(data_manager_mock, resource_manager_mock):
ingestor = IngestorManager(
kafka_servers='localhost:9092',
redis_data={'host': 'localhost', 'port': 6379},
lease_ttl=60,
heartbeat_ttl=60,
poll_interval=5,
mongo_connection_string='mongodb://localhost:27017',
mongo_database='sientia',
export_to_kafka=False,
logger=MagicMock(),
notification_handler=MagicMock(),
metadata=metadata['metadata'],
metrics_controller=MagicMock(),
)
ingestor.send_notification = MagicMock()
ingestor.send_notification_async = AsyncMock()
ingestor.emit_metric = AsyncMock()
return ingestor
@patch('ingestor.managers.ingestor_manager.OpcManager')
@patch('ingestor.managers.ingestor_manager.DataManager')
@patch('ingestor.managers.ingestor_manager.ResourceManager')
@patch('ingestor.managers.ingestor_manager.NotificationHandler')
def test___init__(
notification_handler_mock, resource_manager_mock, data_manager_mock, opc_manager_mock
):
ingestor = IngestorManager(
kafka_servers='localhost:9092',
redis_data={'host': 'localhost', 'port': 6379},
lease_ttl=60,
heartbeat_ttl=60,
poll_interval=5,
mongo_connection_string='mongodb://localhost:27017',
mongo_database='sientia',
export_to_kafka=False,
logger=MagicMock(),
notification_handler=MagicMock(),
metadata=metadata['metadata'],
metrics_controller=MagicMock(),
)
opc_manager_mock.assert_not_called()
data_manager_mock.assert_called_once_with(
kafka_servers='localhost:9092',
mongo_connection_string='mongodb://localhost:27017',
mongo_database='sientia',
export_to_kafka=False,
metadata=metadata['metadata'],
logger=ingestor.logger,
notification_handler=ingestor.notification_handler,
metrics_controller=ingestor.metrics_controller,
)
resource_manager_mock.assert_called_once_with(
host='localhost',
port=6379,
lease_ttl=60,
heartbeat_ttl=60,
metadata=metadata['metadata'],
logger=ingestor.logger,
notification_handler=ingestor.notification_handler,
username=None,
password=None,
metrics_controller=ingestor.metrics_controller,
)
assert ingestor.poll_interval == 5
assert ingestor.managed_tags == {}
assert ingestor.opc_servers == {}
assert ingestor.opc_managers == {}
assert ingestor.data_manager == data_manager_mock.return_value
assert ingestor.resource_manager == resource_manager_mock.return_value
@mark.asyncio
@patch('ingestor.managers.ingestor_manager.OpcManager')
async def test_initialize_opc_from_config(opc_manager, ingestor_manager):
server_config = {
'name': 'server1',
'url': 'opc.tcp://localhost:4840',
'subscription_period_ms': 1000,
'server_uri': 'http://opcua-server.simulator',
'cert_path': '/path/to/cert',
'private_key_path': '/path/to/private_key',
'server_cert_path': '/path/to/server_cert',
'pod_id': 'test_pod',
}
opc_manager.return_value = MagicMock(connect=AsyncMock())
result = await ingestor_manager.initialize_opc_from_config(server_config)
opc_manager.assert_called_once_with(
name=server_config['name'],
url=server_config['url'],
data_manager=ingestor_manager.data_manager,
logger=ingestor_manager.logger,
subscription_period_ms=server_config['subscription_period_ms'],
server_uri=server_config['server_uri'],
notification_handler=ingestor_manager.notification_handler,
cert_path=server_config['cert_path'],
private_key_path=server_config['private_key_path'],
server_cert_path=server_config['server_cert_path'],
metadata=metadata['metadata'],
metrics_controller=ingestor_manager.metrics_controller,
)
assert result == opc_manager.return_value
result.connect.assert_called_once()
@mark.asyncio
@patch('ingestor.managers.ingestor_manager.OpcManager')
@patch('ingestor.managers.ingestor_manager.traceback')
async def test_initialize_opc_from_config_exception(traceback_mock, opc_manager, ingestor_manager):
server_config = {
'name': 'server1',
'url': 'opc.tcp://localhost:4840',
'subscription_period_ms': 1000,
'server_uri': 'http://opcua-server.simulator',
'cert_path': '/path/to/cert',
'private_key_path': '/path/to/private_key',
'server_cert_path': '/path/to/server_cert',
}
ingestor_manager.logger.error = MagicMock()
opc_manager.side_effect = Exception('Initialization error')
result = await ingestor_manager.initialize_opc_from_config(server_config)
assert result is None
traceback_mock.format_exc.assert_called_once()
ingestor_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id=f'OPC_CONNECTION_ERROR_{server_config["name"]}',
message='Error initializing OPC manager: Initialization error',
block='opc_manager',
level=NotificationLevel.ERROR,
attachment_content=traceback_mock.format_exc.return_value,
)
@mark.asyncio
async def test_shutdown(ingestor_manager):
ingestor_manager.opc_managers = {'server1': AsyncMock(), 'server2': AsyncMock()}
ingestor_manager.data_manager.shutdown = MagicMock()
await ingestor_manager.shutdown()
ingestor_manager.opc_managers['server1'].shutdown.assert_called_once()
ingestor_manager.opc_managers['server2'].shutdown.assert_called_once()
ingestor_manager.data_manager.shutdown.assert_called_once()
@patch('ingestor.managers.ingestor_manager.asyncio')
def test___del__(asyncio_mock, ingestor_manager):
ingestor_manager.shutdown = MagicMock()
ingestor_manager.__del__()
asyncio_mock.run.assert_called_once_with(ingestor_manager.shutdown.return_value)
@mark.asyncio
async def test_remove_server(ingestor_manager):
server1 = AsyncMock()
server2 = AsyncMock()
ingestor_manager.opc_managers = {'server1': server1, 'server2': server2}
ingestor_manager.managed_tags = {
'slot1': {'server1': {'config': 'config1'}, 'server2': {'config': 'config2'}}
}
await ingestor_manager.remove_server('server1')
server1.shutdown.assert_called_once()
server2.shutdown.assert_not_called()
assert 'server1' not in ingestor_manager.opc_managers
assert 'server2' in ingestor_manager.opc_managers
assert ingestor_manager.managed_tags == {'slot1': {'server2': {'config': 'config2'}}}
@mark.asyncio
async def test_remove_server_not_found(ingestor_manager):
server1 = AsyncMock()
server2 = AsyncMock()
ingestor_manager.opc_managers = {'server1': server1, 'server2': server2}
ingestor_manager.managed_tags = {
'slot1': {'server1': {'config': 'config1'}, 'server2': {'config': 'config2'}}
}
await ingestor_manager.remove_server('server3')
server1.shutdown.assert_not_called()
server2.shutdown.assert_not_called()
assert 'server1' in ingestor_manager.opc_managers
assert 'server2' in ingestor_manager.opc_managers
assert ingestor_manager.managed_tags == {
'slot1': {'server1': {'config': 'config1'}, 'server2': {'config': 'config2'}}
}
@mark.asyncio
@patch('ingestor.managers.ingestor_manager.OpcManager')
@patch('ingestor.managers.ingestor_manager.metrics')
async def test_update_opc_servers(metrics, opc_manager, ingestor_manager):
manager1 = MagicMock(config={'config': 'config1'})
manager2 = MagicMock(config={'config': 'config2'})
manager3 = MagicMock(config={'config': 'config3'})
async def mock_initialize_from_config(config):
if config == {'config': 'config1'}:
return manager1
elif config == {'config': 'config2'}:
return manager2
elif config == {'config': 'config3'}:
return manager3
else:
return None
ingestor_manager.initialize_opc_from_config = AsyncMock(side_effect=mock_initialize_from_config)
ingestor_manager.remove_server = AsyncMock()
ingestor_manager.managed_tags = {
'slot1': {
'server1': {'config': 'config1'},
'server2': {'config': 'config2'},
'server5': {'config': 'config5'},
},
'slot2': {
'server3': {'config': 'config3'},
'server1': {'config': 'config1'},
},
}
mock = AsyncMock(config={'config': 'old_config2'})
ingestor_manager.opc_managers['server3'] = AsyncMock(config={'config': 'config3'})
ingestor_manager.opc_managers['server2'] = mock
ingestor_manager.opc_managers['server4'] = AsyncMock()
await ingestor_manager.update_opc_servers()
ingestor_manager.initialize_opc_from_config.assert_any_call({'config': 'config1'})
ingestor_manager.initialize_opc_from_config.assert_any_call({'config': 'config2'})
ingestor_manager.initialize_opc_from_config.assert_any_call({'config': 'config5'})
assert ingestor_manager.initialize_opc_from_config.call_count == 3
assert ingestor_manager.opc_managers['server1'].config == {'config': 'config1'}
assert ingestor_manager.opc_managers['server2'].config == {'config': 'config2'}
assert ingestor_manager.opc_managers['server3'].config == {'config': 'config3'}
ingestor_manager.remove_server.assert_any_call('server4')
ingestor_manager.remove_server.assert_any_call('server5')
assert ingestor_manager.remove_server.call_count == 2
assert ingestor_manager.opc_managers['server2'] != mock
ingestor_manager.emit_metric.assert_called_once_with(
metric_object=metrics.OPC_MANAGERS_ACTIVE,
method='set',
value=len(ingestor_manager.opc_managers),
tags={'pod_id': ingestor_manager.pod_id},
)
@mark.asyncio
@patch('ingestor.managers.ingestor_manager.metrics')
async def test_check_opc_servers_integrity_all_healthy(metrics, ingestor_manager):
# Setup mock OPC managers
opc_manager1 = MagicMock()
opc_manager1.check_cycles = AsyncMock(return_value=None)
opc_manager1.check_opc_listenning = AsyncMock(return_value=False)
opc_manager1.config = {'config': 'config1'}
opc_manager2 = MagicMock()
opc_manager2.check_cycles = AsyncMock(return_value=None)
opc_manager2.check_opc_listenning = AsyncMock(return_value=False)
opc_manager2.config = {'config': 'config2'}
ingestor_manager.opc_managers = {'server1': opc_manager1, 'server2': opc_manager2}
# Mock the initialize_opc_from_config method
ingestor_manager.initialize_opc_from_config = MagicMock()
# Call the method
await ingestor_manager.check_opc_servers_integrity()
# Verify that check_cycles and check_opc_listenning were called for each server
opc_manager1.check_cycles.assert_called_once()
opc_manager1.check_opc_listenning.assert_called_once()
opc_manager2.check_cycles.assert_called_once()
opc_manager2.check_opc_listenning.assert_called_once()
# Verify that no reinitialization was needed
ingestor_manager.initialize_opc_from_config.assert_not_called()
ingestor_manager.emit_metric.assert_called_once_with(
metric_object=metrics.OPC_MANAGERS_ACTIVE,
method='set',
value=len(ingestor_manager.opc_managers),
tags={'pod_id': ingestor_manager.pod_id},
)
@mark.asyncio
async def test_check_opc_servers_integrity_server_lost(ingestor_manager):
# Setup mock OPC manager that will be lost
opc_manager = AsyncMock()
opc_manager.check_cycles.return_value = None
opc_manager.check_opc_listenning.return_value = True # Server is lost
opc_manager.config = {'config': 'config1'}
ingestor_manager.opc_managers = {'server1': opc_manager}
ingestor_manager.managed_tags = {'slot1': MagicMock()}
# Call the method
await ingestor_manager.check_opc_servers_integrity()
assert 'server1' not in ingestor_manager.opc_managers
ingestor_manager.managed_tags['slot1'].pop.assert_called_once_with('server1', None)
def test_check_opc_servers_integrity_server_lost_with_tags(ingestor_manager):
# Setup mock OPC manager that will be lost
opc_manager = MagicMock()
opc_manager.check_cycles.return_value = None
opc_manager.check_opc_listenning.return_value = True # Server is lost
opc_manager.config = {'config': 'config1'}
ingestor_manager.opc_managers = {'server1': opc_manager}
# Setup managed tags
ingestor_manager.managed_tags = {
'slot1': {'server1': {'config': 'config1', 'tags': {'tag1': 'value1'}}}
}
# Mock the initialize_opc_from_config method to return a new manager
new_manager = MagicMock()
ingestor_manager.initialize_opc_from_config = MagicMock(return_value=new_manager)
# Call the method
ingestor_manager.check_opc_servers_integrity()
@mark.asyncio
async def test_declare_active(ingestor_manager):
ingestor_manager.resource_manager.ingestor_heartbeat = AsyncMock()
await ingestor_manager.declare_active()
ingestor_manager.resource_manager.ingestor_heartbeat.assert_called_once()
@mark.asyncio
async def test_get_active_ingestors(ingestor_manager):
ingestor_manager.resource_manager.get_all_ingestors = AsyncMock()
await ingestor_manager.get_active_ingestors()
ingestor_manager.resource_manager.get_all_ingestors.assert_called_once()
@mark.asyncio
async def test_get_active_ingestors_empty(ingestor_manager):
ingestor_manager.resource_manager.get_all_ingestors = AsyncMock(return_value=None)
result = await ingestor_manager.get_active_ingestors()
assert result == []
ingestor_manager.resource_manager.get_all_ingestors.assert_called_once()
@patch('ingestor.managers.ingestor_manager.metrics')
@mark.asyncio
async def test_get_number_of_leases_success(metrics, ingestor_manager):
ingestor_manager.resource_manager.get_all_leases = AsyncMock(return_value=['lease1', 'lease2'])
result = await ingestor_manager.get_number_of_leases()
assert result == 2
ingestor_manager.emit_metric.assert_called_once_with(
metric_object=metrics.LEASES_TOTAL,
method='set',
value=2,
tags={'pod_id': ingestor_manager.pod_id},
)
ingestor_manager.resource_manager.get_all_leases.assert_called_once()
@patch('ingestor.managers.ingestor_manager.metrics')
@mark.asyncio
async def test_get_number_of_slots_success(metrics, ingestor_manager):
ingestor_manager.resource_manager.get_all_slots = AsyncMock(return_value=['slot1', 'slot2'])
result = await ingestor_manager.get_number_of_slots()
assert result == 2
ingestor_manager.resource_manager.get_all_slots.assert_called_once()
ingestor_manager.emit_metric.assert_called_once_with(
metric_object=metrics.SLOTS_TOTAL,
method='set',
value=2,
tags={'pod_id': ingestor_manager.pod_id},
)
@patch('ingestor.managers.ingestor_manager.metrics')
@mark.asyncio
async def test_get_number_of_slots_empty(metrics, ingestor_manager):
ingestor_manager.resource_manager.get_all_slots = AsyncMock(return_value=None)
result = await ingestor_manager.get_number_of_slots()
assert result == 0
ingestor_manager.resource_manager.get_all_slots.assert_called_once()
ingestor_manager.emit_metric.assert_called_once_with(
metric_object=metrics.SLOTS_TOTAL,
method='set',
value=0,
tags={'pod_id': ingestor_manager.pod_id},
)
@mark.asyncio
async def test_get_slot_leases_1_success(ingestor_manager):
ingestor_manager.resource_manager.lease_tag = AsyncMock(return_value=True)
ingestor_manager.resource_manager.get_tag_slot = AsyncMock(return_value={'tags': ['tag1']})
ingestor_manager.number_of_slots = 1
result = await ingestor_manager.get_slot_leases()
assert result == {'1': {'tags': ['tag1']}}
@mark.asyncio
async def test_get_slot_leases_2_success(ingestor_manager):
ingestor_manager.resource_manager.lease_tag = AsyncMock(side_effect=[True, True])
ingestor_manager.resource_manager.get_tag_slot = AsyncMock(
side_effect=[{'tags': ['tag1']}, {'tags': ['tag2']}]
)
ingestor_manager.number_of_slots = 2
result = await ingestor_manager.get_slot_leases(max_slots=2)
assert result == {'1': {'tags': ['tag1']}, '2': {'tags': ['tag2']}}
@mark.asyncio
async def test_get_slot_leases_2_1_none(ingestor_manager):
ingestor_manager.resource_manager.lease_tag = AsyncMock(side_effect=[True, True])
ingestor_manager.resource_manager.get_tag_slot = AsyncMock(
side_effect=[None, {'tags': ['tag1']}]
)
ingestor_manager.number_of_slots = 1
result = await ingestor_manager.get_slot_leases(max_slots=1)
assert result == {}
@mark.asyncio
async def test_get_slot_leases_1_failure(ingestor_manager):
ingestor_manager.resource_manager.lease_tag = AsyncMock(return_value=False)
ingestor_manager.resource_manager.get_tag_slot = AsyncMock(return_value={'tags': ['tag1']})
result = await ingestor_manager.get_slot_leases()
ingestor_manager.resource_manager.get_tag_slot.assert_not_called()
assert result == {}
@mark.asyncio
async def test_unsubscribe_slot(ingestor_manager):
ingestor_manager.managed_tags = {
'slot1': {'server1': {'tags': 'config1'}, 'server2': {'tags': 'config2'}},
'slot2': {'server3': {'tags': 'config3'}, 'server1': {'tags': 'config1'}},
}
ingestor_manager.opc_managers = {
'server1': AsyncMock(),
'server2': AsyncMock(),
'server3': AsyncMock(),
}
await ingestor_manager.unsubscribe_slot('slot1')
ingestor_manager.opc_managers['server1'].unsubscribe.assert_called_once_with('slot1')
ingestor_manager.opc_managers['server2'].unsubscribe.assert_called_once_with('slot1')
ingestor_manager.opc_managers['server3'].unsubscribe.assert_not_called()
@mark.asyncio
async def test_update_slot_config(ingestor_manager):
ingestor_manager.managed_tags = {
'slot1': {'config': 'old_config'},
'slot2': {'config': 'new_config'},
'slot3': {'config': 'old_config'},
}
ingestor_manager.resource_manager.get_tag_slot = AsyncMock(
side_effect=[{'config': 'updated_config'}, {'config': 'new_config'}, None]
)
ingestor_manager.resource_manager.renew_tag_lease = AsyncMock()
await ingestor_manager.update_slot_config()
assert ingestor_manager.managed_tags['slot1'] == {'config': 'updated_config'}
assert ingestor_manager.managed_tags['slot2'] == {'config': 'new_config'}
assert 'slot3' not in ingestor_manager.managed_tags
ingestor_manager.resource_manager.renew_tag_lease.assert_any_call('slot1')
ingestor_manager.resource_manager.renew_tag_lease.assert_any_call('slot3')
ingestor_manager.resource_manager.renew_tag_lease.assert_any_call('slot2')
assert ingestor_manager.resource_manager.renew_tag_lease.call_count == 3
@patch('ingestor.managers.ingestor_manager.metrics')
@mark.asyncio
async def test_drop_slot_leases(metrics, ingestor_manager):
ingestor_manager.resource_manager.drop_tag_lease = AsyncMock()
await ingestor_manager.drop_slot_leases(['1', '2'])
ingestor_manager.resource_manager.drop_tag_lease.assert_any_call('1')
ingestor_manager.resource_manager.drop_tag_lease.assert_any_call('2')
ingestor_manager.emit_metric.assert_has_calls(
[
call(
metric_object=metrics.SLOTS_RELEASED,
method='inc',
value=1,
tags={'pod_id': ingestor_manager.pod_id},
),
call(
metric_object=metrics.SLOTS_RELEASED,
method='inc',
value=1,
tags={'pod_id': ingestor_manager.pod_id},
),
]
)
@mark.asyncio
async def test_manage_server_no_server(ingestor_manager):
ingestor_manager.opc_managers = {'server1': MagicMock(), 'server2': MagicMock()}
server_config = {'tags': 'config1'}
result = await ingestor_manager.manage_server('slot1', 'server3', server_config, server_config)
assert result == 1
ingestor_manager.opc_managers['server1'].create_subscription.assert_not_called()
ingestor_manager.opc_managers['server1'].subscribe.assert_not_called()
@mark.asyncio
async def test_manage_server_create_subscription_failure(ingestor_manager):
ingestor_manager.opc_managers = {'server1': MagicMock(), 'server2': MagicMock()}
ingestor_manager.subscriptions = {'server1': MagicMock()}
server_config = {'tags': 'config1'}
ingestor_manager.opc_managers['server1'].create_subscription.side_effect = Exception(
'Subscription error'
)
result = await ingestor_manager.manage_server('slot1', 'server1', server_config, server_config)
assert result == 2
ingestor_manager.opc_managers['server1'].create_subscription.assert_called_once_with('slot1')
ingestor_manager.opc_managers['server1'].subscribe.assert_not_called()
@mark.asyncio
async def test_manage_server(ingestor_manager):
ingestor_manager.opc_managers = {'server1': AsyncMock(), 'server2': AsyncMock()}
ingestor_manager.subscriptions = {'server1': AsyncMock()}
server_config = {'tags': 'config1'}
result = await ingestor_manager.manage_server('slot1', 'server1', server_config, server_config)
assert result == 0
ingestor_manager.opc_managers['server1'].create_subscription.assert_called_once_with('slot1')
ingestor_manager.opc_managers['server1'].subscribe.assert_called_once_with(
'slot1', 'config1', ingestor_manager.poll_interval
)
@patch('ingestor.managers.ingestor_manager.traceback')
@mark.asyncio
async def test_manage_server_subscribe_failure(traceback_mock, ingestor_manager):
ingestor_manager.opc_managers = {'server1': AsyncMock(), 'server2': AsyncMock()}
ingestor_manager.subscriptions = {'server1': AsyncMock()}
server_config = {'tags': 'config1'}
ingestor_manager.opc_managers['server1'].subscribe.side_effect = Exception('Subscription error')
result = await ingestor_manager.manage_server('slot1', 'server1', server_config, server_config)
assert result == 2
ingestor_manager.opc_managers['server1'].create_subscription.assert_called_once_with('slot1')
ingestor_manager.opc_managers['server1'].subscribe.assert_called_once_with(
'slot1', 'config1', ingestor_manager.poll_interval
)
ingestor_manager.opc_managers['server1'].unsubscribe.assert_called_once_with('slot1')
traceback_mock.format_exc.assert_called_once()
ingestor_manager.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='OPC_SUBSCRIPTION_ERROR_slot1:server1',
message="Failed to subscribe to tags from slot1:server1\n{'tags': 'config1'}: Subscription error",
block='opc_manager',
level=NotificationLevel.ERROR,
attachment_content=traceback_mock.format_exc.return_value,
)
ingestor_manager.logger.warning.assert_any_call(
'Removing subscription from server server1 for slot slot1'
)
@mark.asyncio
async def test_subscribe_to_tags(ingestor_manager):
ingestor_manager.manage_server = AsyncMock(side_effect=[0, 1, 2])
ingestor_manager.managed_tags = {'slot1': MagicMock(), 'slot2': MagicMock()}
ingestor_manager.opc_managers = {
'server1': AsyncMock(),
'server2': AsyncMock(),
'server3': AsyncMock(),
}
ingestor_manager.subscriptions = {'server1': AsyncMock()}
tags = {
'slot1': {
'server1': {'tags': 'config1'},
'server2': {'tags': 'config2'},
'server3': {'tags': 'config3'},
}
}
await ingestor_manager.subscribe_to_tags(tags)
ingestor_manager.manage_server.assert_any_call('slot1', 'server1', {'tags': 'config1'}, tags)
ingestor_manager.manage_server.assert_any_call('slot1', 'server2', {'tags': 'config2'}, tags)
ingestor_manager.manage_server.assert_any_call('slot1', 'server3', {'tags': 'config3'}, tags)
assert ingestor_manager.manage_server.call_count == 3
ingestor_manager.managed_tags['slot1'].pop.assert_called_once_with('server3', None)

View File

@@ -0,0 +1,627 @@
import json
from datetime import datetime
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
import pytest
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from ingestor.managers.opc_manager import OpcManager
tags = {
'ns=3;i=1001': {
'aggregation_function': 'LTS',
'frequency': 1000,
'max_value': 100,
'min_value': 0,
'tag_name': 'Counter',
},
'ns=3;i=1003': {
'aggregation_function': 'AVG',
'frequency': 1000,
'max_value': 100,
'min_value': 0,
'tag_name': 'Random',
},
'ns=3;i=1004': {
'aggregation_function': 'MDN',
'frequency': 1000,
'max_value': 100,
'min_value': 0,
'tag_name': 'Sawtooth',
},
}
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@fixture
@patch('ingestor.managers.opc_manager.metrics')
def raw_opc_manager(mock_metrics):
opc_manager = OpcManager(
name='TestConnector',
url='opc.tcp://localhost:4840',
data_manager=AsyncMock(),
subscription_period_ms=1000,
logger=MagicMock(),
server_uri='opc.tcp://localhost:4840',
notification_handler=MagicMock(),
metadata=metadata['metadata'],
metrics_controller=AsyncMock(),
)
opc_manager.emit_metric = AsyncMock()
opc_manager.send_notification_async = AsyncMock()
opc_manager.send_notification = MagicMock()
return opc_manager
@fixture
def opc_manager(raw_opc_manager):
raw_opc_manager.client = AsyncMock()
raw_opc_manager.cert_path = 'cert.pem'
raw_opc_manager.private_key_path = 'private_key.pem'
raw_opc_manager.server_cert_path = 'server_cert.pem'
return raw_opc_manager
@fixture
def opc_manager_subscribed(opc_manager):
opc_manager.subscriptions['sub1'] = AsyncMock()
return opc_manager
def test___str__(opc_manager):
assert (
str(opc_manager)
== 'OpcManager(name=TestConnector, url=opc.tcp://localhost:4840, server_uri=opc.tcp://localhost:4840)\nnodes={}, subscriptions={}'
)
@mark.asyncio
async def test_shutdown_success(opc_manager):
opc_manager.disconnect = AsyncMock()
await opc_manager.shutdown()
opc_manager.disconnect.assert_called_once()
@mark.asyncio
async def test_shutdown_error(opc_manager):
opc_manager.disconnect = AsyncMock(side_effect=Exception('Test error'))
await opc_manager.shutdown()
opc_manager.logger.error.assert_called_once_with('Error during cleanup: Test error')
@mark.asyncio
async def test_set_security_success(opc_manager):
await opc_manager.set_security()
assert opc_manager.client.application_uri == opc_manager.server_uri
opc_manager.client.set_security.assert_called_once_with(
SecurityPolicyBasic256,
certificate=opc_manager.cert_path,
private_key=opc_manager.private_key_path,
server_certificate=opc_manager.server_cert_path,
)
assert opc_manager.client.secure_channel_timeout == 10000000
assert opc_manager.client.session_timeout == 10000000
@mark.asyncio
async def test_set_security_no_cert(opc_manager):
opc_manager.cert_path = None
opc_manager.private_key_path = None
try:
await opc_manager.set_security()
except ValueError as e:
assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
else:
raise AssertionError('ValueError not raised')
assert opc_manager.client.set_security.call_count == 0
@mark.asyncio
@patch('ingestor.managers.opc_manager.metrics')
@patch('ingestor.managers.opc_manager.Client')
async def test_connect_no_security(client, mock_metrics, raw_opc_manager):
raw_opc_manager.set_security = AsyncMock()
client.return_value = AsyncMock()
await raw_opc_manager.connect()
client.assert_called_once_with(raw_opc_manager.url, timeout=10, watchdog_intervall=3600000)
raw_opc_manager.client.connect.assert_called_once()
raw_opc_manager.set_security.assert_not_called()
raw_opc_manager.emit_metric.assert_has_calls(
[
call(
metric_object=mock_metrics.OPC_CONNECTIONS_TOTAL,
tags={
'pod_id': raw_opc_manager.pod_id,
'server_name': raw_opc_manager.name,
},
),
call(
metric_object=mock_metrics.OPC_CONNECTION_STATUS,
method='set',
value=1,
tags={
'pod_id': raw_opc_manager.pod_id,
'server_name': raw_opc_manager.name,
'server_url': raw_opc_manager.url,
},
),
],
any_order=True,
)
@mark.asyncio
@patch('ingestor.managers.opc_manager.Client')
async def test_connect_with_security(client, raw_opc_manager):
raw_opc_manager.cert_path = 'cert.pem'
raw_opc_manager.private_key_path = 'private_key.pem'
raw_opc_manager.server_cert_path = 'server_cert.pem'
raw_opc_manager.set_security = AsyncMock()
client.return_value = AsyncMock()
await raw_opc_manager.connect()
client.assert_called_once_with(raw_opc_manager.url, timeout=10, watchdog_intervall=3600000)
raw_opc_manager.client.connect.assert_called_once()
raw_opc_manager.set_security.assert_called_once()
@mark.asyncio
@patch('ingestor.managers.opc_manager.Client')
@patch('ingestor.managers.opc_manager.metrics')
async def test_connect_exception_handling_and_metrics(
mock_metrics_module, mock_opc_client_class, raw_opc_manager
):
mock_client_instance = mock_opc_client_class.return_value
simulated_error_message = 'Erro de conexão simulado'
mock_client_instance.connect.side_effect = Exception(simulated_error_message)
raw_opc_manager.disconnect = AsyncMock()
opc_manager_instance = raw_opc_manager
opc_manager_instance.cert_path = None
with pytest.raises(Exception, match=simulated_error_message):
await opc_manager_instance.connect()
opc_manager_instance.disconnect.assert_called_once()
@mark.asyncio
async def test_create_subscription_no_client(raw_opc_manager):
try:
await raw_opc_manager.create_subscription('sub1')
except ValueError as e:
assert str(e) == 'Client not connected. Call connect first.'
else:
raise AssertionError('ValueError not raised')
@mark.asyncio
async def test_create_subscription_success_has_period(opc_manager):
await opc_manager.create_subscription('sub1')
opc_manager.client.create_subscription.assert_called_once_with(
opc_manager.subscription_period_ms, opc_manager
)
assert opc_manager.subscriptions['sub1'] is not None
@mark.asyncio
async def test_create_subscription_success_no_period(opc_manager):
await opc_manager.create_subscription('sub1')
opc_manager.client.create_subscription.assert_called_once_with(
opc_manager.subscription_period_ms, opc_manager
)
assert opc_manager.subscriptions['sub1'] is not None
@patch('ingestor.managers.opc_manager.metrics')
@mark.asyncio
async def test_create_subscription_with_metrics(metrics, opc_manager):
await opc_manager.create_subscription('sub1')
opc_manager.emit_metric.assert_called_once_with(
metric_object=metrics.OPC_SUBSCRIPTIONS_CREATED,
method='inc',
value=1,
tags={
'pod_id': opc_manager.pod_id,
'server_name': opc_manager.name,
'slot_name': 'sub1',
},
)
@patch('ingestor.managers.opc_manager.metrics')
@mark.asyncio
async def test_create_subscription_exception_during_client_call(
mock_metrics_module, raw_opc_manager
):
opc_manager_instance = raw_opc_manager
opc_manager_instance.client = AsyncMock()
subscription_name = 'test_sub_client_error'
simulated_error_message = 'Falha ao criar subscrição no cliente OPC'
opc_manager_instance.client.create_subscription.side_effect = Exception(simulated_error_message)
with pytest.raises(Exception, match=simulated_error_message):
await opc_manager_instance.create_subscription(subscription_name)
opc_manager_instance.client.create_subscription.assert_called_once_with(
opc_manager_instance.subscription_period_ms, opc_manager_instance
)
opc_manager_instance.logger.error.assert_called_once_with(
f'Failed to create subscription {subscription_name} on {opc_manager_instance.name}: {simulated_error_message}'
)
mock_metrics_module.OPC_SUBSCRIPTIONS_CREATED.labels.assert_not_called()
@patch('ingestor.managers.opc_manager.metrics')
@mark.asyncio
async def test_subscribe_no_subscription(metrics, opc_manager):
try:
await opc_manager.subscribe('sub1', tags, 1000)
except ValueError as e:
assert str(e) == 'Subscription not created. Call create_subscription first.'
else:
raise AssertionError('ValueError not raised')
metrics.OPC_TAGS_SUBSCRIBED.labels.assert_not_called()
@mark.asyncio
async def test_subscribe_success(opc_manager_subscribed):
opc_manager_subscribed.client.get_node = MagicMock()
opc_manager_subscribed.nodes = {'ns=3;i=1001': 'data'}
await opc_manager_subscribed.subscribe('sub1', tags, 1000)
assert opc_manager_subscribed.nodes == tags
opc_manager_subscribed.subscriptions['sub1'].subscribe_data_change.assert_called_once_with(
[opc_manager_subscribed.client.get_node(n) for n in tags]
)
@mark.asyncio
async def test_unsubscribe_no_subscription(opc_manager):
await opc_manager.unsubscribe('sub1')
opc_manager.logger.warning.assert_called_once_with(
"Subscription 'sub1' not found. Cannot unsubscribe."
)
assert opc_manager.subscriptions.get('sub1') is None
@mark.asyncio
async def test_unsubscribe_success(opc_manager_subscribed):
await opc_manager_subscribed.unsubscribe('sub1')
assert opc_manager_subscribed.subscriptions.get('sub1') is None
@mark.asyncio
async def test_disconnection_fallback_success(opc_manager):
opc_manager.client = AsyncMock()
opc_manager.client.disconnect.return_value = True
result = await opc_manager.disconnection_fallback()
assert result == []
@mark.asyncio
async def test_disconnection_fallback_fail(opc_manager):
opc_manager.client = AsyncMock()
opc_manager.client.disconnect.side_effect = Exception('Test error')
result = await opc_manager.disconnection_fallback()
assert result == [
{'attempt': 1, 'error': 'Test error', 'traceback': ANY},
{'attempt': 2, 'error': 'Test error', 'traceback': ANY},
{'attempt': 3, 'error': 'Test error', 'traceback': ANY},
{'attempt': 4, 'error': 'Test error', 'traceback': ANY},
{'attempt': 5, 'error': 'Test error', 'traceback': ANY},
]
assert opc_manager.client.disconnect.call_count == 5
@mark.asyncio
async def test_disconnect_success(opc_manager_subscribed):
opc_manager_subscribed.client = MagicMock()
await opc_manager_subscribed.disconnect()
opc_manager_subscribed.subscriptions['sub1'].delete.assert_called_once()
assert opc_manager_subscribed.client is None
@mark.asyncio
async def test_disconnect_no_client(opc_manager_subscribed):
opc_manager_subscribed.client = None
assert await opc_manager_subscribed.disconnect() is None
opc_manager_subscribed.logger.warning.assert_has_calls(
[
call('Client already disconnected.'),
]
)
@mark.asyncio
async def test_disconnect_error_unsubscribe(opc_manager_subscribed):
opc_manager_subscribed.client = MagicMock(disconnect=AsyncMock())
opc_manager_subscribed.subscriptions['sub1'] = MagicMock(
delete=AsyncMock(side_effect=Exception('Test error'))
)
await opc_manager_subscribed.disconnect()
opc_manager_subscribed.subscriptions['sub1'].delete.assert_called_once()
opc_manager_subscribed.client = None
opc_manager_subscribed.logger.error.assert_called_once_with(
'Failed to clean up subscription: Test error'
)
@mark.asyncio
async def test_disconnect_error(opc_manager_subscribed):
opc_manager_subscribed.client = MagicMock()
opc_manager_subscribed.client.disconnect = MagicMock(side_effect=Exception('Test error'))
await opc_manager_subscribed.disconnect()
opc_manager_subscribed.subscriptions['sub1'].delete.assert_called_once()
opc_manager_subscribed.client = None
@patch('ingestor.managers.opc_manager.metrics')
@mark.asyncio
async def test_disconnect_metrics_on_successful_path(mock_metrics_module, raw_opc_manager):
mock_metrics_module.OPC_CONNECTION_STATUS.reset_mock()
mock_metrics_module.OPC_TAGS_SUBSCRIBED.reset_mock()
raw_opc_manager.client = MagicMock()
mock_sub1 = MagicMock()
mock_sub2 = MagicMock()
raw_opc_manager.subscriptions = {'sub1': mock_sub1, 'sub2': mock_sub2}
await raw_opc_manager.disconnect()
raw_opc_manager.emit_metric.assert_has_calls(
[
call(
metric_object=mock_metrics_module.OPC_CONNECTION_STATUS,
method='set',
value=0,
tags={
'pod_id': raw_opc_manager.pod_id,
'server_name': raw_opc_manager.name,
'server_url': raw_opc_manager.url,
},
),
call(
metric_object=mock_metrics_module.OPC_TAGS_SUBSCRIBED,
method='set',
value=0,
tags={
'pod_id': raw_opc_manager.pod_id,
'server_name': raw_opc_manager.name,
},
),
],
any_order=True,
)
@patch('ingestor.managers.opc_manager.metrics')
@mark.asyncio
async def test_datachange_notification(metrics, opc_manager_subscribed):
data = MagicMock(
monitored_item=MagicMock(
Value=MagicMock(
Value=MagicMock(Value=42),
SourceTimestamp=datetime.strptime('2021-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
)
)
)
opc_manager_subscribed.nodes = {
'ns=3;i=1001': {
'tag_name': 'Counter',
'cycle_rule': {'cycle_increment': 1.0, 'cycle_count': 2},
'topics': ['topic1', 'topic2'],
}
}
await opc_manager_subscribed.datachange_notification('ns=3;i=1001', None, data)
opc_manager_subscribed.data_manager.publish.assert_any_call(
'topic1',
{
'tag': 'ns=3;i=1001',
'name': 'Counter',
'timestamp': '2021-01-01 00:00:00-0300',
'value': 42,
},
)
opc_manager_subscribed.data_manager.publish.assert_any_call(
'topic2',
{
'tag': 'ns=3;i=1001',
'name': 'Counter',
'timestamp': '2021-01-01 00:00:00-0300',
'value': 42,
},
)
assert opc_manager_subscribed.nodes['ns=3;i=1001']['cycle_rule']['cycle_count'] == 0
opc_manager_subscribed.emit_metric.assert_called_once_with(
metric_object=metrics.OPC_CYCLES_WITHOUT_DATA,
method='set',
value=0,
tags={
'pod_id': opc_manager_subscribed.pod_id,
'server_name': opc_manager_subscribed.name,
},
)
@mark.asyncio
async def test_check_cycles_no_notification(opc_manager):
# Setup: node with cycle_count just below threshold
opc_manager.nodes = {
'ns=3;i=1001': {
'tag_name': 'Counter',
'cycle_rule': {'cycle_increment': 1.0, 'cycle_count': 3.0},
}
}
await opc_manager.check_cycles()
# After one increment, cycle_count = 4.0, still below threshold
assert opc_manager.nodes['ns=3;i=1001']['cycle_rule']['cycle_count'] == pytest.approx(4.0)
opc_manager.send_notification_async.assert_not_called()
@mark.asyncio
async def test_check_cycles_triggers_notification(opc_manager):
# Setup: node with cycle_count just below threshold, increment will cross threshold
opc_manager.nodes = {
'ns=3;i=1001': {
'tag_name': 'Counter',
'cycle_rule': {'cycle_increment': 2.5, 'cycle_count': 3.0},
}
}
opc_manager.notification_handler.build_and_send_notification = MagicMock()
await opc_manager.check_cycles()
# After increment, cycle_count = 5.5, should trigger notification
assert opc_manager.nodes['ns=3;i=1001']['cycle_rule']['cycle_count'] == pytest.approx(5.5)
opc_manager.send_notification_async.assert_called_once_with(
notification_id='TAG_ns=3;i=1001:Counter_LISTENNING_STOPPED',
message='5.5 cycles without receive from ns=3;i=1001:Counter',
block='opc_manager',
level=NotificationLevel.WARNING,
metadata=metadata['metadata'],
)
@patch('ingestor.managers.opc_manager.metrics')
@mark.asyncio
async def test_check_opc_listenning_no_notification(metrics, opc_manager):
opc_manager.non_receive_count = 3
result = await opc_manager.check_opc_listenning()
assert opc_manager.non_receive_count == 4
opc_manager.send_notification_async.assert_not_called()
assert result is False
opc_manager.emit_metric.assert_called_once_with(
metric_object=metrics.OPC_CYCLES_WITHOUT_DATA,
method='set',
value=opc_manager.non_receive_count,
tags={
'pod_id': opc_manager.pod_id,
'server_name': opc_manager.name,
},
)
@mark.asyncio
async def test_check_opc_listenning_warning_notification(opc_manager):
opc_manager.non_receive_count = 4
result = await opc_manager.check_opc_listenning()
assert opc_manager.non_receive_count == 5
opc_manager.send_notification_async.assert_called_once_with(
notification_id=f'OPC_LISTENNING_STOPPED__{opc_manager.name}',
message=f'5 cycles without receive from OPC {opc_manager.name}. Tags: {json.dumps(opc_manager.nodes)}',
block='opc_manager',
level=NotificationLevel.ERROR,
metadata=metadata['metadata'],
)
assert result is False
@patch('ingestor.managers.opc_manager.metrics')
@mark.asyncio
async def test_check_opc_listenning_error_notification_and_retry(metrics, opc_manager):
opc_manager.non_receive_count = 14
result = await opc_manager.check_opc_listenning()
assert opc_manager.non_receive_count == 15
# Should be called twice: once for 5, once for 15
assert opc_manager.send_notification_async.call_count == 2
calls = opc_manager.send_notification_async.call_args_list
# First call: 5 cycles warning
assert calls[0].kwargs == {
'notification_id': f'OPC_LISTENNING_STOPPED__{opc_manager.name}',
'message': f'15 cycles without receive from OPC {opc_manager.name}. Tags: {json.dumps(opc_manager.nodes)}',
'block': 'opc_manager',
'level': NotificationLevel.ERROR,
'metadata': metadata['metadata'],
}
# Second call: 15 cycles retry
assert calls[1].kwargs == {
'notification_id': f'OPC_CONNECTION_RETRY__{opc_manager.name}',
'message': f'Retrying to connect to server {opc_manager.name}',
'block': 'opc_manager',
'level': NotificationLevel.ERROR,
'metadata': metadata['metadata'],
}
assert result is True
opc_manager.emit_metric.assert_has_calls(
[
call(
metric_object=metrics.OPC_CYCLES_WITHOUT_DATA,
method='set',
value=opc_manager.non_receive_count,
tags={
'pod_id': opc_manager.pod_id,
'server_name': opc_manager.name,
},
),
]
)
opc_manager.emit_metric.assert_has_calls(
[
call(
metric_object=metrics.OPC_RECONNECTIONS_TOTAL,
tags={
'pod_id': opc_manager.pod_id,
'server_name': opc_manager.name,
},
),
]
)

View File

@@ -0,0 +1,178 @@
from unittest.mock import AsyncMock, MagicMock, patch
from pytest import fixture, mark, raises
from ingestor.managers.resource_manager import ResourceManager
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@patch('ingestor.managers.resource_manager.RedisRepository')
def test___init__(redis_repository):
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = MagicMock()
resource_manager = ResourceManager(
host='localhost',
port=6379,
lease_ttl=10,
heartbeat_ttl=10,
metadata=metadata['metadata'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
assert resource_manager.redis_repository == redis_repository.return_value
redis_repository.assert_called_once_with(
host='localhost',
port=6379,
username=None,
password=None,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
redis_repository.return_value.redis_client.ping.assert_called_once()
@patch('ingestor.managers.resource_manager.RedisRepository')
def test___init__connection_failure(redis_repository):
redis_repository.return_value.redis_client.ping.side_effect = Exception('Connection failed')
with raises(Exception, match='Connection failed'):
ResourceManager(
host='localhost',
port=6379,
lease_ttl=10,
heartbeat_ttl=10,
metadata=metadata['metadata'],
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
redis_repository.return_value.redis_client.ping.assert_called_once()
redis_repository.logger.error.assert_called_once_with(
'Failed to connect to Redis: Connection failed'
)
redis_repository.return_value.redis_client.ping.assert_called_once()
@fixture
@patch('ingestor.managers.resource_manager.RedisRepository')
def resource_manager(redis_repository):
resource_manager = ResourceManager(
host='localhost',
port=6379,
lease_ttl=10,
heartbeat_ttl=10,
metadata=metadata['metadata'],
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
resource_manager.send_notification = MagicMock()
resource_manager.send_notification_async = AsyncMock()
resource_manager.emit_metric = AsyncMock()
resource_manager.redis_repository = AsyncMock()
return resource_manager
@mark.asyncio
async def test_get_tag_slot(resource_manager):
result = await resource_manager.get_tag_slot('id')
assert result == resource_manager.redis_repository.get.return_value
resource_manager.redis_repository.get.assert_called_once_with(
'slot:opc_tags:id', metadata=metadata['metadata']
)
@mark.asyncio
async def test_ingestor_heartbeat(resource_manager):
await resource_manager.ingestor_heartbeat()
resource_manager.redis_repository.set.assert_called_once_with(
'heartbeat:ingestor:localhost', 1, ttl=10, metadata=metadata['metadata']
)
@mark.asyncio
async def test_lease_tag(resource_manager):
output = await resource_manager.lease_tag('tag_id')
assert output is resource_manager.redis_repository.set.return_value
resource_manager.redis_repository.set.assert_called_once_with(
'lease:opc_tags:tag_id', 'localhost', ttl=10, nx=True, metadata=metadata['metadata']
)
@mark.asyncio
async def test_renew_tag_lease_success(resource_manager):
resource_manager.redis_repository.get.return_value = 'localhost'
resource_manager.redis_repository.expire.return_value = True
result = await resource_manager.renew_tag_lease('tag_id')
assert result is resource_manager.redis_repository.expire.return_value
resource_manager.redis_repository.get.assert_called_once_with(
'lease:opc_tags:tag_id', metadata=metadata['metadata']
)
resource_manager.redis_repository.expire.assert_called_once_with(
'lease:opc_tags:tag_id', 10, metadata=metadata['metadata']
)
@mark.asyncio
async def test_renew_tag_lease_failure(resource_manager):
resource_manager.redis_repository.get.return_value = 'other_pod_id'
result = await resource_manager.renew_tag_lease('tag_id')
assert result is False
resource_manager.redis_repository.get.assert_called_once_with(
'lease:opc_tags:tag_id', metadata=metadata['metadata']
)
resource_manager.redis_repository.expire.assert_not_called()
@mark.asyncio
async def test_drop_tag_lease(resource_manager):
await resource_manager.drop_tag_lease('tag_id')
resource_manager.redis_repository.delete.assert_called_once_with(
'lease:opc_tags:tag_id', metadata=metadata['metadata']
)
@mark.asyncio
async def test_get_all_ingestors(resource_manager):
resource_manager.redis_repository.keys.return_value = ['ingestor1', 'ingestor2']
result = await resource_manager.get_all_ingestors()
assert result == ['ingestor1', 'ingestor2']
resource_manager.redis_repository.keys.assert_called_once_with(
'heartbeat:ingestor:*', metadata=metadata['metadata']
)
@mark.asyncio
async def test_get_all_slots(resource_manager):
resource_manager.redis_repository.keys.return_value = ['slot1', 'slot2']
result = await resource_manager.get_all_slots()
assert result == ['slot1', 'slot2']
resource_manager.redis_repository.keys.assert_called_once_with(
'slot:opc_tags:*', metadata=metadata['metadata']
)
@mark.asyncio
async def test_get_all_leases(resource_manager):
resource_manager.redis_repository.keys.return_value = ['lease1', 'lease2']
result = await resource_manager.get_all_leases()
assert result == ['lease1', 'lease2']
resource_manager.redis_repository.keys.assert_called_once_with(
'lease:opc_tags:*', metadata=metadata['metadata']
)