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

import json
from unittest.mock import MagicMock, patch
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.redis.Redis')
def redis_activity(_mock_redis_client):
    logger = MagicMock()
    notification_handler = MagicMock(spec=NotificationHandler)
    return Redis(host='localhost', port=6379,
                 logger=logger, notification_handler=notification_handler,
                 username='test', password='test')


@patch('scouter.activities.redis.redis.Redis')
def test_redis_initialization(mock_redis_client):
    """Test Redis activity initialization"""
    redis_activity = Redis(host='localhost', port=6379,
                           logger=MagicMock(), notification_handler=MagicMock(),
                           username='test', password='test')
    assert redis_activity.host == 'localhost'
    assert redis_activity.port == 6379
    assert redis_activity.username == 'test'
    assert redis_activity.password == 'test'
    mock_redis_client.assert_called_once_with(
        host='localhost',
        port=6379,
        decode_responses=True,
        username='test',
        password='test'
    )


def test_get_existing_key(redis_activity):
    """Test getting an existing key from Redis"""
    test_data = {'key': 'value'}
    redis_activity.redis_client.get.return_value = json.dumps(test_data)

    result = redis_activity.get('test_key')

    assert result == test_data
    redis_activity.redis_client.get.assert_called_once_with('test_key')


def test_get_nonexistent_key(redis_activity):
    """Test getting a non-existent key from Redis"""
    redis_activity.redis_client.get.return_value = None

    result = redis_activity.get('nonexistent_key')

    assert result is None
    redis_activity.redis_client.get.assert_called_once_with('nonexistent_key')


def test_set_key(redis_activity):
    """Test setting a key in Redis"""
    test_data = {'key': 'value'}

    redis_activity.set('test_key', test_data, ttl=300)

    redis_activity.redis_client.set.assert_called_once_with(
        'test_key',
        json.dumps(test_data),
        ex=300
    )


@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 = {
        '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] == '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 = {
        '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] == '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 = {
        '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 = {
        '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 == {}