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

View File

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

View File

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

View File

@@ -1,5 +1,5 @@
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 pandas as pd
@@ -21,6 +21,9 @@ def gates_fixture():
metrics_controller=metrics_controller,
)
gates.send_notification = MagicMock()
gates.send_notification_async = AsyncMock()
gates.emit_metric = AsyncMock()
gates.logger = logger
gates.notification_handler = notification_handler
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
async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture):
"""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
assert len(result['tag']) == 2
assert 'tag2' not in result['tag']
gates_fixture.send_notification.assert_called_once()
gates_fixture.send_notification_async.assert_called_once()
@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
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
@@ -134,7 +144,7 @@ async def test_data_quality_gate_with_multiple_filters(gates_fixture):
'timestamp': {0: '2023-01-01', 3: '2023-01-04'},
}
# 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
@@ -182,8 +192,8 @@ async def test_data_quality_gate_with_filter_error(gates_fixture):
# Verify error notification is sent and data is unchanged
assert len(result['tag']) == 1
gates_fixture.send_notification.assert_called_once()
call_args = gates_fixture.send_notification.call_args[1]
gates_fixture.send_notification_async.assert_called_once()
call_args = gates_fixture.send_notification_async.call_args[1]
assert call_args['notification_id'] == 'DATA_QUALITY_GATE_ISSUES'
assert call_args['level'] == NotificationLevel.ERROR
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'),
],
)
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."""
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
# Check notification was sent for invalid function
if aggr_function == 'invalid':
gates_fixture.send_notification.assert_called_once()
gates_fixture.send_notification_async.assert_called_once()
else:
gates_fixture.send_notification.assert_not_called()
gates_fixture.send_notification_async.assert_not_called()
@pytest.mark.asyncio
@@ -297,7 +308,7 @@ async def test_aggregate_data(gates_fixture):
@pytest.mark.asyncio
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 = {
'data': [
@@ -324,7 +335,7 @@ async def test_aggregate_data_with_continue(gates_fixture):
# Verify
assert result == expected_result
gates_fixture.send_notification.assert_not_called()
gates_fixture.send_notification_async.assert_not_called()
@pytest.mark.asyncio
@@ -352,7 +363,7 @@ async def test_aggregate_data_raise_exception(gates_fixture):
await gates_fixture.aggregate_data(input_data)
except Exception as e:
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'],
notification_id='AGGREGATION_ISSUES',
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"""
mongodb_activity.mongodb_repository.find.side_effect = Exception('test')
mongodb_activity.send_notification = MagicMock()
mongodb_activity.send_notification_async = AsyncMock()
mongodb_activity.emit_metric = AsyncMock()
try:
await mongodb_activity.load_latest_data(
@@ -165,7 +167,7 @@ async def test_load_latest_data_error(mongodb_activity):
except Exception as e:
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'},
notification_id='MONGO_LOAD_ERROR',
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')
def test_redis_initialization(mock_redis_repository):
"""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_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')
try:
@@ -121,7 +129,7 @@ async def test_get_last_data_timestamp_error(redis_activity):
except Exception as e:
assert str(e) == 'test'
redis_activity.send_notification.assert_called_once_with(
redis_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR',
message='Error getting last data timestamp: test',
@@ -198,8 +206,8 @@ async def test_put_last_data_timestamp_error(redis_activity):
).to_dict('records'),
}
redis_activity.send_notification = MagicMock()
redis_activity.redis_repository.set = MagicMock(side_effect=Exception('test'))
redis_activity.send_notification_async = AsyncMock()
redis_activity.redis_repository.set = AsyncMock(side_effect=Exception('test'))
try:
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:
assert str(e) == 'test'
redis_activity.send_notification.assert_called_once_with(
redis_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='REDIS_SET_ERROR',
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.send_notification = MagicMock()
redis_activity.send_notification_async = AsyncMock()
try:
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:
assert str(e) == 'test'
redis_activity.send_notification.assert_called_once_with(
redis_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR',
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
redis_activity.redis_repository.get = AsyncMock(return_value=existing_data)
redis_activity.redis_repository.set = AsyncMock(side_effect=Exception('test'))
redis_activity.send_notification = MagicMock()
redis_activity.send_notification_async = AsyncMock()
try:
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):
"""Test store_data_package error"""
redis_activity.redis_repository.set = AsyncMock(side_effect=ValueError('test'))
redis_activity.send_notification = MagicMock()
redis_activity.send_notification_async = AsyncMock()
test_data = {
**metadata,
@@ -511,7 +519,7 @@ async def test_store_data_package_error(redis_activity):
with pytest.raises(ValueError):
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'],
notification_id='REDIS_SET_ERROR',
message='Error setting data package: test',