sientia-dataops-scouter_temporal/tests/activities/test_redis.py

from unittest.mock import MagicMock, patch, ANY
from datetime import datetime
import pytest
import numpy as np
from pandas import DataFrame
from sientia_do.notifications.handlers import NotificationHandler
from scouter.activities.redis import Redis


@pytest.fixture
@patch('scouter.activities.redis.RedisBase.__init__')
def redis_activity(_mock_redis_init):
    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
    return activity


@patch('scouter.activities.redis.RedisBase.__init__')
def test_redis_initialization(mock_redis_init):
    """Test Redis activity initialization"""
    logger = MagicMock()
    notification_handler = MagicMock(spec=NotificationHandler)
    Redis(host='localhost', port=6379,
          logger=logger, notification_handler=notification_handler,
          username='test', password='test')
    mock_redis_init.assert_called_once_with(
        ANY,
        'localhost',
        6379,
        'test',
        'test',
        logger,
        notification_handler
    )


metadata = {
    'metadata': {
        'model_id': 'test_model_id',
        'model_name': 'test_model',
        'schedule_name': 'test_schedule',
        'workflow_name': 'scouter'
    }
}


@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')
    }

    # Mock get to return None for new key
    redis_activity.get = MagicMock(return_value=None)
    redis_activity.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.set.assert_called_once()
    args, kwargs = redis_activity.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(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')
    }

    # Mock get to return existing data
    redis_activity.get = MagicMock(return_value=existing_data)
    redis_activity.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'},
        'variable': {0: 'sensor1', 1: 'sensor2', 2: 'sensor3'},
        'value': {0: 25.5, 1: 28.0, 2: 42.0},
        'model_id': {0: 1, 1: 1, 2: 1}
    }
    assert result == expected_result

    # Verify set was called with correct arguments
    redis_activity.set.assert_called_once()
    args, kwargs = redis_activity.set.call_args
    assert args[0] == 'held_data_test_workflow_test_schedule'
    assert args[1] == {
        'sensor1': 25.5,
        'sensor2': 28.0,
        'sensor3': 42.0,
        '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')
    }

    # Mock get to return None for new key
    redis_activity.get = MagicMock(return_value=None)
    redis_activity.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')
    }

    redis_activity.get = MagicMock(return_value=None)

    # Call the method
    result = await redis_activity.group_and_hold_data(test_data)

    assert result == {}