4351289111
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.
530 lines
16 KiB
Python
530 lines
16 KiB
Python
from datetime import datetime
|
|
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from pandas import DataFrame
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
|
|
from scouter.activities.redis import Redis
|
|
|
|
|
|
@pytest.fixture
|
|
@patch('scouter.activities.redis.RedisRepository')
|
|
def redis_activity(mock_redis_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock(spec=NotificationHandler)
|
|
metrics_controller = MagicMock()
|
|
activity = Redis(
|
|
host='localhost',
|
|
port=6379,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
username='test',
|
|
password='test',
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
activity.redis_client = MagicMock()
|
|
activity.logger = logger
|
|
activity.notification_handler = notification_handler
|
|
activity.pod_id = 'test_pod_id'
|
|
return activity
|
|
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model_id',
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test_schedule',
|
|
'workflow_name': 'scouter',
|
|
}
|
|
}
|
|
|
|
|
|
@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"""
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock(spec=NotificationHandler)
|
|
metrics_controller = MagicMock()
|
|
activity = Redis(
|
|
host='localhost',
|
|
port=6379,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
username='test',
|
|
password='test',
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
mock_redis_repository.assert_called_once_with(
|
|
host='localhost',
|
|
port=6379,
|
|
username='test',
|
|
password='test',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
assert activity.redis_repository is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_last_data_timestamp_none(redis_activity):
|
|
"""Test get_last_data_timestamp"""
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_pipeline',
|
|
'schedule_name': 'test_schedule',
|
|
}
|
|
|
|
redis_activity.redis_repository.get = AsyncMock(return_value=None)
|
|
|
|
result = await redis_activity.get_last_data_timestamp(test_data)
|
|
|
|
redis_activity.redis_repository.get.assert_called_once_with(
|
|
'last_data_timestamp:test_pipeline:test_schedule'
|
|
)
|
|
|
|
assert result is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_last_data_timestamp_not_none(redis_activity):
|
|
"""Test get_last_data_timestamp"""
|
|
test_data = {**metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}
|
|
|
|
redis_activity.redis_repository.get = AsyncMock(return_value='2023-01-01 12:00:00')
|
|
|
|
result = await redis_activity.get_last_data_timestamp(test_data)
|
|
|
|
redis_activity.redis_repository.get.assert_called_once_with(
|
|
'last_data_timestamp:test_pipeline:test_schedule'
|
|
)
|
|
|
|
assert result == '2023-01-01 12:00:00'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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_async = AsyncMock()
|
|
redis_activity.redis_repository.get.side_effect = Exception('test')
|
|
|
|
try:
|
|
await redis_activity.get_last_data_timestamp(test_data)
|
|
|
|
except Exception as e:
|
|
assert str(e) == 'test'
|
|
|
|
redis_activity.send_notification_async.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='REDIS_GET_ERROR',
|
|
message='Error getting last data timestamp: test',
|
|
block='get_last_data_timestamp',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
else:
|
|
raise AssertionError('Expected exception')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_put_last_data_timestamp_empty_dataframe(redis_activity):
|
|
"""Test put_last_data_timestamp with empty dataframe"""
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_pipeline',
|
|
'schedule_name': 'test_schedule',
|
|
'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'),
|
|
}
|
|
|
|
redis_activity.redis_repository.set = MagicMock()
|
|
|
|
result = await redis_activity.put_last_data_timestamp(test_data)
|
|
|
|
assert result is None
|
|
|
|
redis_activity.redis_repository.set.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_put_last_data_timestamp_not_empty_dataframe(redis_activity):
|
|
"""Test put_last_data_timestamp with not empty dataframe"""
|
|
|
|
data = DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor2'],
|
|
'value': [25.5, 30.0],
|
|
'inserted_at': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'],
|
|
}
|
|
)
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_pipeline',
|
|
'schedule_name': 'test_schedule',
|
|
'data': data.to_dict('records'),
|
|
}
|
|
|
|
redis_activity.redis_repository.set = AsyncMock()
|
|
|
|
result = await redis_activity.put_last_data_timestamp(test_data)
|
|
|
|
assert result == '2023-01-01 12:00:01'
|
|
|
|
redis_activity.redis_repository.set.assert_called_once_with(
|
|
'last_data_timestamp:test_pipeline:test_schedule', '2023-01-01 12:00:01', ttl=18000
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_put_last_data_timestamp_error(redis_activity):
|
|
"""Test put_last_data_timestamp error"""
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_pipeline',
|
|
'schedule_name': 'test_schedule',
|
|
'data': DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor2'],
|
|
'value': [25.5, 30.0],
|
|
'inserted_at': ['2023-01-01 12:00:00'] * 2,
|
|
}
|
|
).to_dict('records'),
|
|
}
|
|
|
|
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)
|
|
|
|
except Exception as e:
|
|
assert str(e) == 'test'
|
|
|
|
redis_activity.send_notification_async.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='REDIS_SET_ERROR',
|
|
message='Error setting last data timestamp: test',
|
|
block='put_last_data_timestamp',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
else:
|
|
raise AssertionError('Expected exception')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_group_and_hold_data_new_key(redis_activity):
|
|
"""Test group_and_hold_data with a new key"""
|
|
# Setup
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_pipeline',
|
|
'schedule_name': 'test_schedule',
|
|
'retention_time': 3600,
|
|
'model_id': 1,
|
|
'data': DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor2'],
|
|
'value': [25.5, 30.0],
|
|
'timestamp': ['2023-01-01 12:00:00'] * 2,
|
|
}
|
|
).to_dict('records'),
|
|
'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'},
|
|
'fill_missing_tags': False,
|
|
}
|
|
|
|
# Mock get to return None for new key
|
|
redis_activity.redis_repository.get = AsyncMock(return_value=None)
|
|
redis_activity.redis_repository.set = AsyncMock()
|
|
|
|
# Call the method
|
|
result = await redis_activity.group_and_hold_data(test_data)
|
|
|
|
# Verify the result
|
|
expected_result = {
|
|
'timestamp': {0: '2023-01-01 12:00:00', 1: '2023-01-01 12:00:00'},
|
|
'variable': {0: 'sensor1', 1: 'sensor2'},
|
|
'value': {0: 25.5, 1: 30.0},
|
|
'model_id': {0: 1, 1: 1},
|
|
}
|
|
assert result == expected_result
|
|
|
|
# Verify set was called with correct arguments
|
|
redis_activity.redis_repository.set.assert_called_once_with(
|
|
'held_data_test_pipeline_test_schedule',
|
|
{'sensor1': 25.5, 'sensor2': 30.0, 'timestamp': '2023-01-01 12:00:00'},
|
|
ttl=3600,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_group_and_hold_data_update_existing_fill_missing(redis_activity):
|
|
"""Test updating existing data with group_and_hold_data"""
|
|
# Setup initial data in Redis
|
|
existing_data = {'sensor1': 20.0, 'sensor2': 28.0, 'timestamp': '2023-01-01 11:00:00'}
|
|
|
|
# New data to update with
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
'retention_time': 3600,
|
|
'model_id': 1,
|
|
'data': DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor3'],
|
|
'value': [25.5, 42.0],
|
|
'timestamp': ['2023-01-01 12:00:00'] * 2,
|
|
}
|
|
).to_dict('records'),
|
|
'model_tags': {
|
|
'sensor1': 'sensor1',
|
|
'sensor2': 'sensor2',
|
|
'sensor3': 'sensor3',
|
|
'sensor4': 'sensor4',
|
|
},
|
|
'fill_missing_tags': True,
|
|
}
|
|
|
|
# Mock get to return existing data
|
|
redis_activity.redis_repository.get = AsyncMock(return_value=existing_data)
|
|
redis_activity.redis_repository.set = AsyncMock()
|
|
|
|
# Call the method
|
|
result = await redis_activity.group_and_hold_data(test_data)
|
|
|
|
# Verify the result
|
|
expected_result = {
|
|
'timestamp': {
|
|
0: '2023-01-01 12:00:00',
|
|
1: '2023-01-01 12:00:00',
|
|
2: '2023-01-01 12:00:00',
|
|
3: '2023-01-01 12:00:00',
|
|
},
|
|
'variable': {0: 'sensor1', 1: 'sensor2', 2: 'sensor3', 3: 'sensor4'},
|
|
'value': {0: 25.5, 1: 28.0, 2: 42.0, 3: None},
|
|
'model_id': {0: 1, 1: 1, 2: 1, 3: 1},
|
|
}
|
|
assert result == expected_result
|
|
|
|
# Verify set was called with correct arguments
|
|
redis_activity.redis_repository.set.assert_called_once_with(
|
|
'held_data_test_workflow_test_schedule',
|
|
{
|
|
'sensor1': 25.5,
|
|
'sensor2': 28.0,
|
|
'sensor3': 42.0,
|
|
'sensor4': None,
|
|
'timestamp': '2023-01-01 12:00:00',
|
|
},
|
|
ttl=3600,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_group_and_hold_data_with_none_values(redis_activity):
|
|
"""Test handling of None values in group_and_hold_data"""
|
|
# Setup test data with None values
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
'retention_time': 3600,
|
|
'model_id': 1,
|
|
'data': DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor2'],
|
|
'value': [None, 30.0],
|
|
'timestamp': [datetime(2023, 1, 1, 12, 0, 0)] * 2,
|
|
}
|
|
).to_dict('records'),
|
|
'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'},
|
|
'fill_missing_tags': False,
|
|
}
|
|
|
|
# Mock get to return None for new key
|
|
redis_activity.redis_repository.get = AsyncMock(return_value=None)
|
|
redis_activity.redis_repository.set = AsyncMock()
|
|
|
|
# Call the method
|
|
result = await redis_activity.group_and_hold_data(test_data)
|
|
|
|
# Verify None was converted to np.nan and values are as expected
|
|
assert np.isnan(result['value'][0])
|
|
assert result['value'][1] == pytest.approx(30.0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_group_and_hold_data_empty_dataframe(redis_activity):
|
|
"""Test group_and_hold_data with empty DataFrame"""
|
|
# Setup test with empty data
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
'retention_time': 3600,
|
|
'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'),
|
|
'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'},
|
|
'fill_missing_tags': False,
|
|
}
|
|
|
|
redis_activity.redis_repository.get = AsyncMock(return_value=None)
|
|
|
|
# Call the method
|
|
result = await redis_activity.group_and_hold_data(test_data)
|
|
|
|
assert result == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_group_and_hold_data_error_get(redis_activity):
|
|
"""Test group_and_hold_data error"""
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
'retention_time': 3600,
|
|
'model_id': 1,
|
|
'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'),
|
|
'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'},
|
|
'fill_missing_tags': False,
|
|
}
|
|
|
|
redis_activity.redis_repository.get = AsyncMock(side_effect=Exception('test'))
|
|
redis_activity.send_notification_async = AsyncMock()
|
|
|
|
try:
|
|
await redis_activity.group_and_hold_data(test_data)
|
|
|
|
except Exception as e:
|
|
assert str(e) == 'test'
|
|
|
|
redis_activity.send_notification_async.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='REDIS_GET_ERROR',
|
|
message='Error getting held data: test',
|
|
block='group_and_hold_data',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
else:
|
|
raise AssertionError('Expected exception')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_group_and_hold_data_error_set(redis_activity):
|
|
"""Test group_and_hold_data error"""
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
'retention_time': 3600,
|
|
'model_id': 1,
|
|
'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'),
|
|
'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'},
|
|
'fill_missing_tags': False,
|
|
}
|
|
|
|
existing_data = {'sensor1': 20.0, 'sensor2': 28.0, 'timestamp': '2023-01-01 11:00:00'}
|
|
|
|
# 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_async = AsyncMock()
|
|
|
|
try:
|
|
await redis_activity.group_and_hold_data(test_data)
|
|
|
|
except Exception as e:
|
|
assert str(e) == 'test'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_data_package(redis_activity):
|
|
"""Test store_data_package"""
|
|
redis_activity.redis_repository.set = AsyncMock()
|
|
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
'held_data': DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor2'],
|
|
'value': [25.5, 30.0],
|
|
'timestamp': ['2023-01-01 12:00:00'] * 2,
|
|
}
|
|
).to_dict(),
|
|
'data': DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor2'],
|
|
'value': [25.5, 30.0],
|
|
'timestamp': ['2023-01-01 12:00:00'] * 2,
|
|
}
|
|
).to_dict(),
|
|
'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'},
|
|
}
|
|
|
|
await redis_activity.store_data_package(test_data)
|
|
|
|
redis_activity.redis_repository.set.assert_called_once_with(
|
|
ANY, {'data': test_data['data'], 'held_data': test_data['held_data']}, ttl=120
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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_async = AsyncMock()
|
|
|
|
test_data = {
|
|
**metadata,
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
'held_data': DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor2'],
|
|
'value': [25.5, 30.0],
|
|
'timestamp': ['2023-01-01 12:00:00'] * 2,
|
|
}
|
|
).to_dict(),
|
|
'data': DataFrame(
|
|
{
|
|
'name': ['sensor1', 'sensor2'],
|
|
'value': [25.5, 30.0],
|
|
'timestamp': ['2023-01-01 12:00:00'] * 2,
|
|
}
|
|
).to_dict(),
|
|
'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'},
|
|
}
|
|
|
|
with pytest.raises(ValueError):
|
|
await redis_activity.store_data_package(test_data)
|
|
|
|
redis_activity.send_notification_async.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='REDIS_SET_ERROR',
|
|
message='Error setting data package: test',
|
|
block='store_data_package',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|