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:
@@ -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,
|
||||
|
||||
@@ -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}',
|
||||
|
||||
@@ -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}',
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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',
|
||||
|
||||
Reference in New Issue
Block a user