0c16075823
Update GITHUB_BRANCH in values.yaml to feature/SIENTIAPDE-1325 for specific external operation metrics. Refactor shutdown methods in Activities class to ensure proper closure of MongoDB, Redis, and Gates connections. Adjust tests to reflect these changes and enhance repository pattern implementation for MongoDB and Redis activities.
521 lines
16 KiB
Python
521 lines
16 KiB
Python
from datetime import datetime
|
|
from unittest.mock import ANY, 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)
|
|
activity = Redis(
|
|
host='localhost',
|
|
port=6379,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
username='test',
|
|
password='test',
|
|
)
|
|
|
|
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.RedisRepository')
|
|
def test_redis_initialization(mock_redis_repository):
|
|
"""Test Redis activity initialization"""
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock(spec=NotificationHandler)
|
|
activity = Redis(
|
|
host='localhost',
|
|
port=6379,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
username='test',
|
|
password='test',
|
|
)
|
|
mock_redis_repository.assert_called_once_with(
|
|
host='localhost',
|
|
port=6379,
|
|
username='test',
|
|
password='test',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
activity.logger = logger
|
|
activity.notification_handler = notification_handler
|
|
activity.pod_id = 'localhost'
|
|
|
|
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.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.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 = MagicMock()
|
|
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.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 = MagicMock()
|
|
|
|
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 = MagicMock()
|
|
redis_activity.redis_repository.set = MagicMock(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.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 = MagicMock(return_value=None)
|
|
redis_activity.redis_repository.set = MagicMock()
|
|
|
|
# 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()
|
|
args, kwargs = redis_activity.redis_repository.set.call_args
|
|
assert args[0] == 'held_data_test_pipeline_test_schedule'
|
|
assert args[1] == {'sensor1': 25.5, 'sensor2': 30.0, 'timestamp': '2023-01-01 12:00:00'}
|
|
assert kwargs['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 = MagicMock(return_value=existing_data)
|
|
redis_activity.redis_repository.set = MagicMock()
|
|
|
|
# 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()
|
|
args, kwargs = redis_activity.redis_repository.set.call_args
|
|
assert args[0] == 'held_data_test_workflow_test_schedule'
|
|
assert args[1] == {
|
|
'sensor1': 25.5,
|
|
'sensor2': 28.0,
|
|
'sensor3': 42.0,
|
|
'sensor4': None,
|
|
'timestamp': '2023-01-01 12:00:00',
|
|
}
|
|
assert kwargs['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 = MagicMock(return_value=None)
|
|
redis_activity.redis_repository.set = MagicMock()
|
|
|
|
# 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 = MagicMock(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 = MagicMock(side_effect=Exception('test'))
|
|
redis_activity.send_notification = MagicMock()
|
|
|
|
try:
|
|
await redis_activity.group_and_hold_data(test_data)
|
|
|
|
except Exception as e:
|
|
assert str(e) == 'test'
|
|
|
|
redis_activity.send_notification.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 = MagicMock(return_value=existing_data)
|
|
redis_activity.redis_repository.set = MagicMock(side_effect=Exception('test'))
|
|
redis_activity.send_notification = MagicMock()
|
|
|
|
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 = MagicMock()
|
|
|
|
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 = MagicMock(side_effect=ValueError('test'))
|
|
redis_activity.send_notification = MagicMock()
|
|
|
|
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.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,
|
|
)
|