SIENTIAPDE-1325

Refactor notification methods in Gates, MongoDB, and Redis classes to use asynchronous send_notification_async. Update related tests to ensure proper async handling and verification of notifications.
This commit is contained in:
vitor-aignosi
2025-11-03 14:57:44 -03:00
parent 8350a3a0c3
commit 4351289111
6 changed files with 57 additions and 36 deletions

View File

@@ -59,7 +59,7 @@ class Gates(SientiaMonitoring):
""" """
SientiaMonitoring.shutdown(self) SientiaMonitoring.shutdown(self)
def apply_aggregation( async def apply_aggregation(
self, values: DataFrame, aggr_function: str, metadata: dict[str, Any] self, values: DataFrame, aggr_function: str, metadata: dict[str, Any]
) -> float | None | str: ) -> float | None | str:
""" """
@@ -106,7 +106,7 @@ class Gates(SientiaMonitoring):
if aggr_function in aggregation_map: if aggr_function in aggregation_map:
return aggregation_map[aggr_function](clean_values) return aggregation_map[aggr_function](clean_values)
else: else:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='AGGREGATION_ISSUES', notification_id='AGGREGATION_ISSUES',
message=f'Invalid aggregation function: {aggr_function}', message=f'Invalid aggregation function: {aggr_function}',
@@ -165,7 +165,7 @@ class Gates(SientiaMonitoring):
# Get the latest timestamp (last row since data is sorted) # Get the latest timestamp (last row since data is sorted)
latest_timestamp = group['timestamp'].iloc[-1] latest_timestamp = group['timestamp'].iloc[-1]
aggr_value = self.apply_aggregation(group, aggr_function, metadata) aggr_value = await self.apply_aggregation(group, aggr_function, metadata)
if aggr_value == 'continue': if aggr_value == 'continue':
continue continue
@@ -197,7 +197,7 @@ class Gates(SientiaMonitoring):
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='AGGREGATION_ISSUES', notification_id='AGGREGATION_ISSUES',
message=f'Error aggregating data: {e}', message=f'Error aggregating data: {e}',
@@ -255,7 +255,7 @@ class Gates(SientiaMonitoring):
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='DATA_QUALITY_GATE_ISSUES', notification_id='DATA_QUALITY_GATE_ISSUES',
message=f'Error applying filter {filter_name}: {e}', message=f'Error applying filter {filter_name}: {e}',
@@ -273,7 +273,7 @@ class Gates(SientiaMonitoring):
message = f'{len(filtered_data)} rows has quality issues: {filter_name}: {policy}' message = f'{len(filtered_data)} rows has quality issues: {filter_name}: {policy}'
attachment = filtered_data.to_string() attachment = filtered_data.to_string()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id=f'DATA_QUALITY_GATE_ISSUES__{filter_name}', notification_id=f'DATA_QUALITY_GATE_ISSUES__{filter_name}',
message=message, message=message,

View File

@@ -143,7 +143,7 @@ class MongoDB(SientiaMonitoring):
return data return data
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='MONGO_LOAD_ERROR', notification_id='MONGO_LOAD_ERROR',
message=f'Error loading data from MongoDB: {e}', message=f'Error loading data from MongoDB: {e}',

View File

@@ -99,7 +99,7 @@ class Redis(SientiaMonitoring):
try: try:
data_hold = await self.redis_repository.get(key) data_hold = await self.redis_repository.get(key)
except Exception as e: except Exception as e:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message=f'Error getting last data timestamp: {e}', message=f'Error getting last data timestamp: {e}',
@@ -157,7 +157,7 @@ class Redis(SientiaMonitoring):
try: try:
await self.redis_repository.set(key, last_data_timestamp, ttl=60 * 60 * 5) await self.redis_repository.set(key, last_data_timestamp, ttl=60 * 60 * 5)
except Exception as e: except Exception as e:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_SET_ERROR', notification_id='REDIS_SET_ERROR',
message=f'Error setting last data timestamp: {e}', message=f'Error setting last data timestamp: {e}',
@@ -210,7 +210,7 @@ class Redis(SientiaMonitoring):
try: try:
data_hold = await self.redis_repository.get(key) data_hold = await self.redis_repository.get(key)
except Exception as e: except Exception as e:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message=f'Error getting held data: {e}', message=f'Error getting held data: {e}',
@@ -262,7 +262,7 @@ class Redis(SientiaMonitoring):
data_hold_melted.reset_index(drop=True, inplace=True) data_hold_melted.reset_index(drop=True, inplace=True)
except Exception as e: except Exception as e:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_SET_ERROR', notification_id='REDIS_SET_ERROR',
message=f'Error setting held data: {e}', message=f'Error setting held data: {e}',
@@ -300,7 +300,7 @@ class Redis(SientiaMonitoring):
try: try:
await self.redis_repository.set(key, cache, ttl=120) await self.redis_repository.set(key, cache, ttl=120)
except Exception as e: except Exception as e:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_SET_ERROR', notification_id='REDIS_SET_ERROR',
message=f'Error setting data package: {e}', message=f'Error setting data package: {e}',

View File

@@ -1,5 +1,5 @@
from typing import Any from typing import Any
from unittest.mock import ANY, MagicMock, Mock, call, patch from unittest.mock import ANY, AsyncMock, MagicMock, Mock, call, patch
import numpy as np import numpy as np
import pandas as pd import pandas as pd
@@ -21,6 +21,9 @@ def gates_fixture():
metrics_controller=metrics_controller, metrics_controller=metrics_controller,
) )
gates.send_notification = MagicMock() gates.send_notification = MagicMock()
gates.send_notification_async = AsyncMock()
gates.emit_metric = AsyncMock()
gates.logger = logger gates.logger = logger
gates.notification_handler = notification_handler gates.notification_handler = notification_handler
gates.metrics_controller = metrics_controller gates.metrics_controller = metrics_controller
@@ -38,6 +41,13 @@ metadata = {
} }
@patch('scouter.activities.gates.SientiaMonitoring')
def test_close(mock_sientia_monitoring, gates_fixture):
"""Test close method."""
gates_fixture.close()
mock_sientia_monitoring.shutdown.assert_called_once()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture): async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture):
"""Test data_quality_gate with NULL_VALUES_FILTER and DISCARD policy.""" """Test data_quality_gate with NULL_VALUES_FILTER and DISCARD policy."""
@@ -64,7 +74,7 @@ async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture):
# Verify # Verify
assert len(result['tag']) == 2 assert len(result['tag']) == 2
assert 'tag2' not in result['tag'] assert 'tag2' not in result['tag']
gates_fixture.send_notification.assert_called_once() gates_fixture.send_notification_async.assert_called_once()
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -97,7 +107,7 @@ async def test_data_quality_gate_with_out_of_bounds_filter_keep(gates_fixture):
# Verify data is kept but notification is sent # Verify data is kept but notification is sent
assert len(result['tag']) == 3 # All rows kept assert len(result['tag']) == 3 # All rows kept
gates_fixture.send_notification.assert_called_once() gates_fixture.send_notification_async.assert_called_once()
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -134,7 +144,7 @@ async def test_data_quality_gate_with_multiple_filters(gates_fixture):
'timestamp': {0: '2023-01-01', 3: '2023-01-04'}, 'timestamp': {0: '2023-01-01', 3: '2023-01-04'},
} }
# Should be called twice (once for each filter) # Should be called twice (once for each filter)
assert gates_fixture.send_notification.call_count == 2 assert gates_fixture.send_notification_async.call_count == 2
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -182,8 +192,8 @@ async def test_data_quality_gate_with_filter_error(gates_fixture):
# Verify error notification is sent and data is unchanged # Verify error notification is sent and data is unchanged
assert len(result['tag']) == 1 assert len(result['tag']) == 1
gates_fixture.send_notification.assert_called_once() gates_fixture.send_notification_async.assert_called_once()
call_args = gates_fixture.send_notification.call_args[1] call_args = gates_fixture.send_notification_async.call_args[1]
assert call_args['notification_id'] == 'DATA_QUALITY_GATE_ISSUES' assert call_args['notification_id'] == 'DATA_QUALITY_GATE_ISSUES'
assert call_args['level'] == NotificationLevel.ERROR assert call_args['level'] == NotificationLevel.ERROR
assert 'Filter error' in call_args['message'] assert 'Filter error' in call_args['message']
@@ -246,16 +256,17 @@ async def test_data_quality_gate_with_no_filters(gates_fixture):
(pd.DataFrame({'value': [1.0, 2.0]}), 'invalid', 'continue'), (pd.DataFrame({'value': [1.0, 2.0]}), 'invalid', 'continue'),
], ],
) )
def test_apply_aggregation(gates_fixture, group_data, aggr_function, expected_result): @pytest.mark.asyncio
async def test_apply_aggregation(gates_fixture, group_data, aggr_function, expected_result):
"""Test apply_aggregation method with various scenarios.""" """Test apply_aggregation method with various scenarios."""
result = gates_fixture.apply_aggregation(group_data, aggr_function, metadata) result = await gates_fixture.apply_aggregation(group_data, aggr_function, metadata)
assert result == expected_result assert result == expected_result
# Check notification was sent for invalid function # Check notification was sent for invalid function
if aggr_function == 'invalid': if aggr_function == 'invalid':
gates_fixture.send_notification.assert_called_once() gates_fixture.send_notification_async.assert_called_once()
else: else:
gates_fixture.send_notification.assert_not_called() gates_fixture.send_notification_async.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -297,7 +308,7 @@ async def test_aggregate_data(gates_fixture):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_aggregate_data_with_continue(gates_fixture): async def test_aggregate_data_with_continue(gates_fixture):
gates_fixture.apply_aggregation = MagicMock(return_value='continue') gates_fixture.apply_aggregation = AsyncMock(return_value='continue')
input_data = { input_data = {
'data': [ 'data': [
@@ -324,7 +335,7 @@ async def test_aggregate_data_with_continue(gates_fixture):
# Verify # Verify
assert result == expected_result assert result == expected_result
gates_fixture.send_notification.assert_not_called() gates_fixture.send_notification_async.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -352,7 +363,7 @@ async def test_aggregate_data_raise_exception(gates_fixture):
await gates_fixture.aggregate_data(input_data) await gates_fixture.aggregate_data(input_data)
except Exception as e: except Exception as e:
assert str(e) == 'Test exception' assert str(e) == 'Test exception'
gates_fixture.send_notification.assert_called_once_with( gates_fixture.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='AGGREGATION_ISSUES', notification_id='AGGREGATION_ISSUES',
message='Error aggregating data: Test exception', message='Error aggregating data: Test exception',

View File

@@ -153,6 +153,8 @@ async def test_load_latest_data_error(mongodb_activity):
"""Test load_latest_data""" """Test load_latest_data"""
mongodb_activity.mongodb_repository.find.side_effect = Exception('test') mongodb_activity.mongodb_repository.find.side_effect = Exception('test')
mongodb_activity.send_notification = MagicMock() mongodb_activity.send_notification = MagicMock()
mongodb_activity.send_notification_async = AsyncMock()
mongodb_activity.emit_metric = AsyncMock()
try: try:
await mongodb_activity.load_latest_data( await mongodb_activity.load_latest_data(
@@ -165,7 +167,7 @@ async def test_load_latest_data_error(mongodb_activity):
except Exception as e: except Exception as e:
assert str(e) == 'test' assert str(e) == 'test'
mongodb_activity.send_notification.assert_called_once_with( mongodb_activity.send_notification_async.assert_called_once_with(
metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
notification_id='MONGO_LOAD_ERROR', notification_id='MONGO_LOAD_ERROR',
message='Error loading data from MongoDB: test', message='Error loading data from MongoDB: test',

View File

@@ -43,6 +43,14 @@ metadata = {
} }
@patch('scouter.activities.redis.SientiaMonitoring')
def test_close(mock_sientia_monitoring, redis_activity):
"""Test close method."""
redis_activity.close()
redis_activity.redis_repository.close.assert_called_once()
mock_sientia_monitoring.shutdown.assert_called_once()
@patch('scouter.activities.redis.RedisRepository') @patch('scouter.activities.redis.RedisRepository')
def test_redis_initialization(mock_redis_repository): def test_redis_initialization(mock_redis_repository):
"""Test Redis activity initialization""" """Test Redis activity initialization"""
@@ -112,7 +120,7 @@ async def test_get_last_data_timestamp_error(redis_activity):
"""Test get_last_data_timestamp error""" """Test get_last_data_timestamp error"""
test_data = {**metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'} test_data = {**metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}
redis_activity.send_notification = MagicMock() redis_activity.send_notification_async = AsyncMock()
redis_activity.redis_repository.get.side_effect = Exception('test') redis_activity.redis_repository.get.side_effect = Exception('test')
try: try:
@@ -121,7 +129,7 @@ async def test_get_last_data_timestamp_error(redis_activity):
except Exception as e: except Exception as e:
assert str(e) == 'test' assert str(e) == 'test'
redis_activity.send_notification.assert_called_once_with( redis_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message='Error getting last data timestamp: test', message='Error getting last data timestamp: test',
@@ -198,8 +206,8 @@ async def test_put_last_data_timestamp_error(redis_activity):
).to_dict('records'), ).to_dict('records'),
} }
redis_activity.send_notification = MagicMock() redis_activity.send_notification_async = AsyncMock()
redis_activity.redis_repository.set = MagicMock(side_effect=Exception('test')) redis_activity.redis_repository.set = AsyncMock(side_effect=Exception('test'))
try: try:
await redis_activity.put_last_data_timestamp(test_data) await redis_activity.put_last_data_timestamp(test_data)
@@ -207,7 +215,7 @@ async def test_put_last_data_timestamp_error(redis_activity):
except Exception as e: except Exception as e:
assert str(e) == 'test' assert str(e) == 'test'
redis_activity.send_notification.assert_called_once_with( redis_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='REDIS_SET_ERROR', notification_id='REDIS_SET_ERROR',
message='Error setting last data timestamp: test', message='Error setting last data timestamp: test',
@@ -399,7 +407,7 @@ async def test_group_and_hold_data_error_get(redis_activity):
} }
redis_activity.redis_repository.get = AsyncMock(side_effect=Exception('test')) redis_activity.redis_repository.get = AsyncMock(side_effect=Exception('test'))
redis_activity.send_notification = MagicMock() redis_activity.send_notification_async = AsyncMock()
try: try:
await redis_activity.group_and_hold_data(test_data) await redis_activity.group_and_hold_data(test_data)
@@ -407,7 +415,7 @@ async def test_group_and_hold_data_error_get(redis_activity):
except Exception as e: except Exception as e:
assert str(e) == 'test' assert str(e) == 'test'
redis_activity.send_notification.assert_called_once_with( redis_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message='Error getting held data: test', message='Error getting held data: test',
@@ -439,7 +447,7 @@ async def test_group_and_hold_data_error_set(redis_activity):
# Mock get to return existing data # Mock get to return existing data
redis_activity.redis_repository.get = AsyncMock(return_value=existing_data) redis_activity.redis_repository.get = AsyncMock(return_value=existing_data)
redis_activity.redis_repository.set = AsyncMock(side_effect=Exception('test')) redis_activity.redis_repository.set = AsyncMock(side_effect=Exception('test'))
redis_activity.send_notification = MagicMock() redis_activity.send_notification_async = AsyncMock()
try: try:
await redis_activity.group_and_hold_data(test_data) await redis_activity.group_and_hold_data(test_data)
@@ -485,7 +493,7 @@ async def test_store_data_package(redis_activity):
async def test_store_data_package_error(redis_activity): async def test_store_data_package_error(redis_activity):
"""Test store_data_package error""" """Test store_data_package error"""
redis_activity.redis_repository.set = AsyncMock(side_effect=ValueError('test')) redis_activity.redis_repository.set = AsyncMock(side_effect=ValueError('test'))
redis_activity.send_notification = MagicMock() redis_activity.send_notification_async = AsyncMock()
test_data = { test_data = {
**metadata, **metadata,
@@ -511,7 +519,7 @@ async def test_store_data_package_error(redis_activity):
with pytest.raises(ValueError): with pytest.raises(ValueError):
await redis_activity.store_data_package(test_data) await redis_activity.store_data_package(test_data)
redis_activity.send_notification.assert_called_once_with( redis_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='REDIS_SET_ERROR', notification_id='REDIS_SET_ERROR',
message='Error setting data package: test', message='Error setting data package: test',