SIENTIAPDE-1318

Refactor project structure and update configurations

- Deleted the empty `__init__.py` file to clean up the project structure.
- Renamed the project from "laborious" to "ingestor" in `pyproject.toml`, updating the description accordingly.
- Improved type hints and conditional checks in the `Ingestor`, `DataManager`, `IngestorManager`, and `OpcManager` classes for better code clarity and type safety.
- Enhanced error handling and assertions in various methods to ensure robustness.
- Updated unit tests to reflect changes in class names and error handling improvements.
This commit is contained in:
vitor-aignosi
2025-10-17 14:54:39 -03:00
parent e2462af31c
commit e19b57d4b1
11 changed files with 73 additions and 65 deletions

View File

View File

@@ -72,12 +72,7 @@ class Ingestor:
"""
kafka_servers = getenv('KAFKA_SERVERS', 'localhost:9092')
export_to_kafka = getenv('EXPORT_TO_KAFKA', 'false')
if export_to_kafka and export_to_kafka == 'true':
export_to_kafka = True
else:
export_to_kafka = False
export_to_kafka: bool = getenv('EXPORT_TO_KAFKA', 'false') == 'true'
self.export_to_kafka = export_to_kafka
self.redis_host = getenv('REDIS_HOST', 'localhost')
@@ -110,7 +105,7 @@ class Ingestor:
'schema_name': 'opc_ingestor',
'pod_id': self.pod_id,
}
self.ingestor_manager = None
self.ingestor_manager: IngestorManager | None = None
async def shutdown(self):
"""
@@ -150,8 +145,9 @@ class Ingestor:
else:
# Subscribe to acquired slots
self.ingestor_manager.update_opc_servers()
await self.ingestor_manager.subscribe_to_tags(acquired)
if self.ingestor_manager:
self.ingestor_manager.update_opc_servers()
await self.ingestor_manager.subscribe_to_tags(acquired)
async def prepare_ingestor(self):
"""
@@ -173,7 +169,7 @@ class Ingestor:
"""
self.ingestor_manager = IngestorManager(
kafka_servers=self.kafka_servers,
kafka_servers=','.join(self.kafka_servers),
redis_data={
'host': self.redis_host,
'port': self.redis_port,
@@ -190,6 +186,7 @@ class Ingestor:
notification_handler=self.notification_handler,
export_to_kafka=self.export_to_kafka,
)
assert self.ingestor_manager is not None
# Declare ingestor active
self.ingestor_manager.declare_active()
@@ -221,7 +218,7 @@ class Ingestor:
- Requests a single slot lease to begin processing
"""
if not self.ingestor_manager.managed_tags and number_of_slots > 0:
if self.ingestor_manager and not self.ingestor_manager.managed_tags and number_of_slots > 0:
# This ingestor is active and has no slots, so we need to try to
# Get slot lease
@@ -250,6 +247,8 @@ class Ingestor:
- Logs the number of available slots when attempting to acquire leases.
- Logs the number of extra slots when releasing leases.
"""
if not self.ingestor_manager:
return
if available_slots > 0 and lacking_ingestors > 0:
# Some ingestors are inactive, so there are "available_slots" slots available
@@ -297,6 +296,9 @@ class Ingestor:
- Updates metrics to reflect current state
"""
if not self.ingestor_manager:
return
self.logger.debug(f'Current managed tags: {self.ingestor_manager.managed_tags}')
await self.ingestor_manager.update_opc_servers()
@@ -361,6 +363,9 @@ class Ingestor:
and OPC servers.
"""
if not self.ingestor_manager:
return
self.ingestor_manager.declare_active()
self.logger.info('Polling for slot updates...')

View File

@@ -129,7 +129,7 @@ class DataManager(BaseActivity):
self.connection_string = mongo_connection_string
self.database = mongo_database
self.mongo_client = MongoClient(self.connection_string)
self.mongo_client: MongoClient = MongoClient(self.connection_string)
self.mongo_client.server_info()
self.metadata = metadata
@@ -176,7 +176,7 @@ class DataManager(BaseActivity):
def __del__(self):
self.shutdown()
def delivery_report(self, msg: str):
def delivery_report(self, msg):
"""
Callback for successful Kafka message delivery reports.
@@ -191,7 +191,7 @@ class DataManager(BaseActivity):
f'Record successfully produced to {msg.topic} [{msg.partition}] at offset {msg.offset}'
)
def delivery_error(self, err: str):
def delivery_error(self, err):
"""
Callback for Kafka message delivery error reports.
@@ -218,7 +218,7 @@ class DataManager(BaseActivity):
Exception: If there is an error during message delivery, it will be handled by the `delivery_error` callback.
"""
if self.export_to_kafka:
if self.export_to_kafka and self.kafka_producer:
try:
self.logger.debug(f'Publishing message to topic {topic}: {data}')
self.kafka_producer.send(topic=topic, value=data).add_callback(

View File

@@ -72,10 +72,10 @@ class IngestorManager(BaseActivity):
notification_handler: NotificationHandler,
export_to_kafka: bool = False,
):
redis_host = redis_data.get('host')
redis_port = redis_data.get('port')
redis_username = redis_data.get('username', None)
redis_password = redis_data.get('password', None)
redis_host: str = redis_data['host']
redis_port: int = int(redis_data['port'])
redis_username: str | None = redis_data.get('username', None)
redis_password: str | None = redis_data.get('password', None)
self.data_manager = DataManager(
kafka_servers=kafka_servers,
@@ -86,7 +86,7 @@ class IngestorManager(BaseActivity):
logger=logger,
notification_handler=notification_handler,
)
self.opc_managers = {}
self.opc_managers: dict = {}
self.resource_manager = ResourceManager(
host=redis_host,
port=redis_port,
@@ -100,8 +100,8 @@ class IngestorManager(BaseActivity):
)
self.number_of_slots = 0
self.poll_interval = poll_interval
self.managed_tags = {}
self.opc_servers = {}
self.managed_tags: dict = {}
self.opc_servers: dict = {}
self.metadata = metadata
@@ -457,8 +457,7 @@ class IngestorManager(BaseActivity):
- Logs informational messages for updated slot configurations.
"""
removed_slots = []
update = {}
removed_slots: list[str] = []
for slot, _slot_config in self.managed_tags.items():
self.resource_manager.renew_tag_lease(slot)
update = self.resource_manager.get_tag_slot(slot)

View File

@@ -64,21 +64,21 @@ class OpcManager(BaseActivity):
server_uri: str,
notification_handler: NotificationHandler,
metadata: dict,
cert_path: str = None,
private_key_path: str = None,
server_cert_path: str = None,
cert_path: str | None = None,
private_key_path: str | None = None,
server_cert_path: str | None = None,
):
self.url = url
self.name = name
self.server_uri = server_uri
self.data_queue = {}
self.data_queue: dict = {}
self.non_receive_count = 0
self.client = None
self.client: Client | None = None
self.cert_path = cert_path
self.private_key_path = private_key_path
self.server_cert_path = server_cert_path
self.nodes = {}
self.subscriptions = {}
self.nodes: dict = {}
self.subscriptions: dict = {}
self.data_manager = data_manager
self.metadata = metadata
@@ -147,20 +147,21 @@ class OpcManager(BaseActivity):
raise ValueError(
'Certificate and private key paths must be provided for secure connection.'
)
cert = Path(self.cert_path)
private_key = Path(self.private_key_path)
cert = Path(self.cert_path) if self.cert_path else None
private_key = Path(self.private_key_path) if self.private_key_path else None
server_cert = Path(self.server_cert_path) if self.server_cert_path else None
await self.client.set_application_uri(self.server_uri)
self.logger.info('Setting security...')
await self.client.set_security(
SecurityPolicyBasic256,
certificate=str(cert),
private_key=str(private_key),
server_certificate=str(server_cert),
)
await self.client.set_secure_channel_timeout(10000000)
await self.client.set_session_timeout(10000000)
if self.client:
await self.client.set_application_uri(self.server_uri)
self.logger.info('Setting security...')
await self.client.set_security(
SecurityPolicyBasic256,
certificate=str(cert),
private_key=str(private_key),
server_certificate=str(server_cert),
)
await self.client.set_secure_channel_timeout(10000000)
await self.client.set_session_timeout(10000000)
async def connect(self):
"""
@@ -187,6 +188,7 @@ class OpcManager(BaseActivity):
metrics.OPC_CONNECTIONS_TOTAL.labels(pod_id=self.pod_id, server_name=self.name).inc()
try:
self.client = Client(self.url, watchdog_intervall=3600000)
assert self.client is not None # Informa ao mypy que client não é None
self.client.name = self.pod_id
if self.cert_path:
await self.set_security()
@@ -268,6 +270,7 @@ class OpcManager(BaseActivity):
self.logger.info(f'Subscribing to {subscription} on {self.name}...')
self.logger.info(f'Subscribing to nodes: {nodes}')
assert self.client is not None # Informa ao mypy que client não é None
addr_nodes = [self.client.get_node(n) for n in nodes]
self.logger.debug(f'Addr nodes: {addr_nodes}')
self.nodes.update(nodes)
@@ -400,7 +403,8 @@ class OpcManager(BaseActivity):
'value': value,
}
_a = [self.data_manager.publish(e, data) for e in self.nodes[tag]['topics']]
for topic in self.nodes[tag]['topics']:
self.data_manager.publish(topic, data)
def check_cycles(self):
"""

View File

@@ -154,7 +154,7 @@ class ResourceManager(BaseActivity):
)
raise
def get(self, key: str) -> dict:
def get(self, key: str) -> dict | None:
"""
Retrieve a value from Redis by its key and return it as a dictionary.
@@ -176,7 +176,7 @@ class ResourceManager(BaseActivity):
history = self._execute_redis_op('get', self.redis.get, key)
return json.loads(history) if history else None
def get_tag_slot(self, id: str) -> dict:
def get_tag_slot(self, tag_id: str) -> dict | None:
"""
Retrieve the tag slot information for a given ID.
@@ -194,7 +194,7 @@ class ResourceManager(BaseActivity):
and delegates to the get() method for the actual Redis operation.
"""
return self.get(f'slot:opc_tags:{id}')
return self.get(f'slot:opc_tags:{tag_id}')
def ingestor_heartbeat(self) -> None:
"""

View File

@@ -3,9 +3,9 @@ requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[project]
name = "laborious"
name = "ingestor"
version = "0.0.0"
description = "Sientia DataOps Laborious - ML Model Orchestration System"
description = "Sientia DataOps Ingestor - OPC Tag Ingestor"
readme = "README.md"
requires-python = ">=3.11"
authors = [

View File

@@ -135,7 +135,7 @@ def test___init___failure_max_attempts(mongo, kafka):
assert logger_mock.info.call_count == 3
else:
assert False, 'Expected NoBrokersAvailable exception was not raised.'
raise AssertionError('Expected NoBrokersAvailable exception was not raised.')
def test_shutdown_has_producer(data_manager):

View File

@@ -127,7 +127,7 @@ async def test_set_security_no_cert(opc_manager):
except ValueError as e:
assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
else:
assert False, 'ValueError not raised'
raise AssertionError('ValueError not raised')
assert opc_manager.client.set_security.call_count == 0
@@ -218,7 +218,7 @@ async def test_create_subscription_no_client(raw_opc_manager):
except ValueError as e:
assert str(e) == 'Client not connected. Call connect first.'
else:
assert False, 'ValueError not raised'
raise AssertionError('ValueError not raised')
@mark.asyncio
@@ -284,7 +284,7 @@ async def test_subscribe_no_subscription(metrics, opc_manager):
except ValueError as e:
assert str(e) == 'Subscription not created. Call create_subscription first.'
else:
assert False, 'ValueError not raised'
raise AssertionError('ValueError not raised')
metrics.OPC_TAGS_SUBSCRIBED.labels.assert_not_called()
@@ -315,7 +315,7 @@ async def test_unsubscribe_no_subscription(opc_manager):
async def test_unsubscribe_success(opc_manager_subscribed):
await opc_manager_subscribed.unsubscribe('sub1')
opc_manager_subscribed.subscriptions.get('sub1') is None
assert opc_manager_subscribed.subscriptions.get('sub1') is None
@mark.asyncio

View File

@@ -10,7 +10,7 @@ from ingestor import app
# Custom exception to catch os._exit calls
class OsExitCalled(Exception):
class OsExitCalledError(Exception):
def __init__(self, code):
super().__init__(f'os._exit({code}) called')
self.code = code
@@ -18,7 +18,7 @@ class OsExitCalled(Exception):
# Helper function for the os_exit mock's side_effect
def raise_os_exit_with_code(exit_code):
raise OsExitCalled(exit_code)
raise OsExitCalledError(exit_code)
@pytest.fixture
@@ -95,7 +95,7 @@ async def test_main_successful_run_one_loop(mock_app_env, capsys):
mock_exit_signal.is_set.side_effect = [False, True]
mock_app_env['time_time'].side_effect = [10.0, 11.5]
with pytest.raises(OsExitCalled) as excinfo:
with pytest.raises(OsExitCalledError) as excinfo:
await app.main()
assert excinfo.value.code == 0
@@ -133,7 +133,7 @@ async def test_main_prometheus_server_fails_to_start(mock_app_env, capsys):
"""Test the scenario where starting the Prometheus server fails."""
mock_app_env['start_http_server'].side_effect = OSError('Port already in use')
with pytest.raises(OsExitCalled) as excinfo:
with pytest.raises(OsExitCalledError) as excinfo:
await app.main()
assert excinfo.value.code == 1
@@ -160,7 +160,7 @@ async def test_main_loop_exception_handling(mock_app_env, capsys):
mock_ingestor_instance.loop.side_effect = Exception('Test loop exception')
mock_app_env['time_time'].side_effect = [10.0, 10.1]
with pytest.raises(OsExitCalled) as excinfo:
with pytest.raises(OsExitCalledError) as excinfo:
await app.main()
assert excinfo.value.code == 0
@@ -194,7 +194,7 @@ async def test_main_keyboard_interrupt_handling(mock_app_env, capsys):
mock_ingestor_instance.loop.side_effect = KeyboardInterrupt()
mock_app_env['time_time'].side_effect = [10.0, 10.1]
with pytest.raises(OsExitCalled) as excinfo:
with pytest.raises(OsExitCalledError) as excinfo:
await app.main()
assert excinfo.value.code == 0
@@ -223,7 +223,7 @@ async def test_main_multiple_loop_iterations(mock_app_env):
mock_exit_signal.is_set.side_effect = [False, False, False, True]
mock_app_env['time_time'].side_effect = [10.0, 10.1, 10.2, 10.3, 10.4, 10.5]
with pytest.raises(OsExitCalled) as excinfo:
with pytest.raises(OsExitCalledError) as excinfo:
await app.main()
assert excinfo.value.code == 0
@@ -254,7 +254,7 @@ async def test_main_pod_id_used_in_metrics(mock_app_env):
mock_exit_signal.is_set.side_effect = [False, True]
mock_app_env['time_time'].side_effect = [10.0, 11.0]
with pytest.raises(OsExitCalled):
with pytest.raises(OsExitCalledError):
await app.main()
app.metrics.APP_UP.labels.assert_any_call(pod_id='test_pod')
@@ -279,7 +279,7 @@ async def test_run_async_main(mock_app_env, capsys):
# Instead of calling run_async_main() which creates a new event loop,
# we test the main() function directly since that's what run_async_main() would call
with pytest.raises(OsExitCalled) as excinfo:
with pytest.raises(OsExitCalledError) as excinfo:
await app.main()
assert excinfo.value.code == 0
@@ -299,7 +299,7 @@ async def test_main_prepare_ingestor_failure(mock_app_env, capsys):
mock_ingestor_instance.prepare_ingestor.side_effect = Exception('Preparation failed')
mock_exit_signal.is_set.side_effect = [False, True]
with pytest.raises(OsExitCalled) as excinfo:
with pytest.raises(OsExitCalledError) as excinfo:
await app.main()
assert excinfo.value.code == 0

View File

@@ -122,7 +122,7 @@ async def test_prepare_ingestor(ingestor_manager_mock, ingestor):
await ingestor.prepare_ingestor()
ingestor_manager_mock.assert_called_once_with(
kafka_servers=ingestor.kafka_servers,
kafka_servers=','.join(ingestor.kafka_servers),
redis_data={
'host': ingestor.redis_host,
'port': ingestor.redis_port,