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

from unittest.mock import Mock, patch, MagicMock, ANY
import numpy as np
import pandas as pd
import pytest
from sientia_do.notifications.models import NotificationLevel
from scouter.activities.gates import Gates


@pytest.fixture
def gates_fixture():
    """Fixture to create a Gates instance with mocked dependencies."""
    logger = Mock()
    notification_handler = MagicMock()
    return Gates(logger=logger, notification_handler=notification_handler)


@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."""
    # Setup test data
    input_data = {
        'filters': {
            'NULL_VALUES_FILTER': 'DISCARD'
        },
        'data': {
            'tag': ['tag1', 'tag2', 'tag3'],
            'value': [1.0, None, 3.0],
            'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03'],
        },
        'model_tags': {
            'tag1': {'data_range': [0, 100]},
            'tag2': {'data_range': [0, 100]},
            'tag3': {'data_range': [0, 100]}
        }
    }

    # Execute
    result = await gates_fixture.data_quality_gate(input_data)

    # Verify
    assert len(result['tag']) == 2
    assert 'tag2' not in result['tag']
    gates_fixture.notification_handler.build_and_send_notification.assert_called_once()


@pytest.mark.asyncio
async def test_data_quality_gate_with_out_of_bounds_filter_keep(gates_fixture):
    """Test data_quality_gate with OUT_OF_BOUNDS_FILTER and KEEP policy."""
    # Setup test data with out of bounds values
    input_data = {
        'filters': {
            'OUT_OF_BOUNDS_FILTER': 'KEEP'
        },
        'data': {
            'tag': ['tag1', 'tag2', 'tag3'],
            'value': [1.0, 200.0, 3.0],
            'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03']
        },
        'model_tags': {
            'tag1': {'data_range': [0, 100]},
            'tag2': {'data_range': [0, 100]},
            'tag3': {'data_range': [0, 100]}
        }
    }

    # Mock the out_of_bounds_filter to return rows with out of bounds values
    with patch('scouter.activities.gates.quality_gate_filters', {
        'OUT_OF_BOUNDS_FILTER': lambda df: df[df['tag'] == 'tag2']
    }):
        # Execute
        result = await gates_fixture.data_quality_gate(input_data)

    # Verify data is kept but notification is sent
    assert len(result['tag']) == 3  # All rows kept
    gates_fixture.notification_handler.build_and_send_notification.assert_called_once()


@pytest.mark.asyncio
async def test_data_quality_gate_with_multiple_filters(gates_fixture):
    """Test data_quality_gate with multiple filters."""
    # Setup test data
    input_data = {
        'filters': {
            'NULL_VALUES_FILTER': 'DISCARD',
            'OUT_OF_BOUNDS_FILTER': 'DISCARD'
        },
        'data': {
            'tag': ['tag1', 'tag2', 'tag3', 'tag4'],
            'name': ['tag1', 'tag2', 'tag3', 'tag4'],
            'value': [1.0, None, 300.0, 4.0],
            'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04']
        },
        'model_tags': {
            'tag1': {'data_range': [0, 100]},
            'tag2': {'data_range': [0, 100]},
            'tag3': {'data_range': [0, 100]},
            'tag4': {'data_range': [0, 100]}
        }
    }

    result = await gates_fixture.data_quality_gate(input_data)

    # Verify only tag1 and tag4 remain (tag2 has null, tag3 is out of bounds)
    assert result == {'tag': {0: 'tag1', 3: 'tag4'}, 'name': {0: 'tag1', 3: 'tag4'}, 'value': {
        0: 1.0, 3: 4.0}, 'timestamp': {0: '2023-01-01', 3: '2023-01-04'}}
    # Should be called twice (once for each filter)
    assert gates_fixture.notification_handler.build_and_send_notification.call_count == 2


@pytest.mark.asyncio
async def test_data_quality_gate_with_unknown_filter(gates_fixture):
    """Test data_quality_gate with an unknown filter."""
    # Setup test data with unknown filter
    input_data = {
        'filters': {
            'UNKNOWN_FILTER': 'DISCARD'
        },
        'data': {
            'tag': ['tag1'],
            'name': ['tag1'],
            'value': [1.0],
            'timestamp': ['2023-01-01']
        },
        'model_tags': {
            'tag1': {'data_range': [0, 100]}
        }
    }

    # Execute
    result = await gates_fixture.data_quality_gate(input_data)

    # Verify data is unchanged and warning is logged
    assert len(result['tag']) == 1
    gates_fixture.logger.warning.assert_called_once_with(
        "Filter UNKNOWN_FILTER not found")


@pytest.mark.asyncio
async def test_data_quality_gate_with_filter_error(gates_fixture):
    """Test data_quality_gate when a filter raises an exception."""
    # Setup test data
    input_data = {
        'filters': {
            'NULL_VALUES_FILTER': 'DISCARD'
        },
        'data': {
            'tag': ['tag1'],
            'name': ['tag1'],
            'value': [1.0],
            'timestamp': ['2023-01-01']
        },
        'model_tags': {
            'tag1': {'data_range': [0, 100]}
        }
    }

    # Mock the filter to raise an exception
    def failing_filter(_, _model_tags):
        raise ValueError("Filter error")

    with patch('scouter.activities.gates.quality_gate_filters', {
        'NULL_VALUES_FILTER': failing_filter
    }):
        # Execute
        result = await gates_fixture.data_quality_gate(input_data)

    # Verify error notification is sent and data is unchanged
    assert len(result['tag']) == 1
    gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
    call_args = gates_fixture.notification_handler.build_and_send_notification.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']


@pytest.mark.asyncio
async def test_data_quality_gate_with_empty_data(gates_fixture):
    """Test data_quality_gate with empty input data."""
    # Setup empty input data
    input_data = {
        'filters': {
            'NULL_VALUES_FILTER': 'DISCARD'
        },
        'data': {
            'tag': [],
            'name': [],
            'value': [],
            'timestamp': []
        },
        'model_tags': {}
    }

    # Execute
    result = await gates_fixture.data_quality_gate(input_data)

    # Verify empty result and no notifications
    assert len(result['tag']) == 0
    gates_fixture.notification_handler.build_and_send_notification.assert_not_called()


@pytest.mark.asyncio
async def test_data_quality_gate_with_no_filters(gates_fixture):
    """Test data_quality_gate with no filters specified."""
    # Setup test data with no filters
    input_data = {
        'filters': {},
        'data': {
            'tag': ['tag1'],
            'name': ['tag1'],
            'value': [1.0],
            'timestamp': ['2023-01-01']
        },
        'model_tags': {
            'tag1': {'data_range': [0, 100]}
        }
    }

    # Execute
    result = await gates_fixture.data_quality_gate(input_data)

    # Verify data is unchanged and no notifications
    assert len(result['tag']) == 1
    gates_fixture.notification_handler.build_and_send_notification.assert_not_called()


@pytest.mark.parametrize(
    "group_data, aggr_function, expected_result",
    [
        # Single value case
        (pd.DataFrame({'value': [10.0]}), 'avg', 10.0),
        # Multiple values with different aggregation functions
        (pd.DataFrame({'value': [1.0, 2.0, 3.0, 4.0]}), 'avg', 2.5),
        (pd.DataFrame({'value': [1.0, 2.0, 3.0, 4.0]}), 'mdn', 2.5),
        (pd.DataFrame({'value': [1.0, 2.0, 3.0, 4.0]}), 'max', 4.0),
        (pd.DataFrame({'value': [1.0, 2.0, 3.0, 4.0]}), 'min', 1.0),
        (pd.DataFrame({'value': [1.0, 2.0, 3.0, 4.0]}), 'lts', 4.0),
        # With NaN values
        (pd.DataFrame({'value': [1.0, np.nan, 3.0, 4.0]}),
         'avg', 2.6666666666666665),
        # Empty group after dropping NaN
        (pd.DataFrame({'value': [np.nan, np.nan]}), 'avg', None),
        # Invalid aggregation function
        (pd.DataFrame({'value': [1.0, 2.0]}), 'invalid', 'continue'),
    ]
)
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)
    assert result == expected_result

    # Check notification was sent for invalid function
    if aggr_function == 'invalid':
        gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
    else:
        gates_fixture.notification_handler.build_and_send_notification.assert_not_called()


@pytest.mark.asyncio
async def test_aggregate_data(gates_fixture):
    """Test aggregate_data method with multiple groups and aggregation functions."""
    input_data = {
        'data': [
            {'tag': 'tag1', 'name': 'name1', 'value': 1.0, 'timestamp': '2023-01-01'},
            {'tag': 'tag1', 'name': 'name1', 'value': 2.0, 'timestamp': '2023-01-02'},
            {'tag': 'tag1', 'name': 'name1', 'value': 3.0, 'timestamp': '2023-01-03'},
            {'tag': 'tag2', 'name': 'name2', 'value': 4.0, 'timestamp': '2023-01-01'},
            {'tag': 'tag2', 'name': 'name2', 'value': 5.0, 'timestamp': '2023-01-02'},
            {'tag': 'tag2', 'name': 'name2', 'value': 6.0, 'timestamp': '2023-01-03'},
            {'tag': 'tag1', 'name': 'name1',
                'value': None, 'timestamp': '2023-01-04'},
        ],
        'model_tags': {
            'name1': {'aggr_function': 'avg'},
            'name2': {'aggr_function': 'max'},
        }
    }

    # Expected result
    expected_result = {'tag': {0: 'tag1', 1: 'tag2'},
                       'name': {0: 'name1', 1: 'name2'},
                       'value': {0: 2.0, 1: 6.0},
                       'timestamp': {0: '2023-01-04', 1: '2023-01-03'},
                       'aggregation_function': {0: 'avg', 1: 'max'}}

    # Execute
    result = await gates_fixture.aggregate_data(input_data)

    # Verify
    assert result == expected_result
    gates_fixture.notification_handler.build_and_send_notification.assert_not_called()


@pytest.mark.asyncio
async def test_aggregate_data_with_continue(gates_fixture):

    gates_fixture.apply_aggregation = MagicMock(return_value='continue')

    input_data = {
        'data': [
            {'tag': 'tag1', 'name': 'name1', 'value': 1.0, 'timestamp': '2023-01-01'},
            {'tag': 'tag1', 'name': 'name1', 'value': 2.0, 'timestamp': '2023-01-02'},
            {'tag': 'tag1', 'name': 'name1', 'value': 3.0, 'timestamp': '2023-01-03'},
            {'tag': 'tag2', 'name': 'name2', 'value': 4.0, 'timestamp': '2023-01-01'},
            {'tag': 'tag2', 'name': 'name2', 'value': 5.0, 'timestamp': '2023-01-02'},
            {'tag': 'tag2', 'name': 'name2', 'value': 6.0, 'timestamp': '2023-01-03'},
            {'tag': 'tag1', 'name': 'name1',
                'value': None, 'timestamp': '2023-01-04'},
        ],
        'model_tags': {
            'name1': {'aggr_function': 'avg'},
            'name2': {'aggr_function': 'max'},
        }
    }

    # Expected result
    expected_result = {}

    # Execute
    result = await gates_fixture.aggregate_data(input_data)

    # Verify
    assert result == expected_result
    gates_fixture.notification_handler.build_and_send_notification.assert_not_called()


@pytest.mark.asyncio
async def test_aggregate_data_raise_exception(gates_fixture):

    gates_fixture.apply_aggregation = MagicMock(
        side_effect=Exception("Test exception"))

    input_data = {
        'data': [
            {'tag': 'tag1', 'name': 'name1', 'value': 1.0, 'timestamp': '2023-01-01'},
            {'tag': 'tag1', 'name': 'name1', 'value': 2.0, 'timestamp': '2023-01-02'},
            {'tag': 'tag1', 'name': 'name1', 'value': 3.0, 'timestamp': '2023-01-03'},
            {'tag': 'tag2', 'name': 'name2', 'value': 4.0, 'timestamp': '2023-01-01'},
            {'tag': 'tag2', 'name': 'name2', 'value': 5.0, 'timestamp': '2023-01-02'},
            {'tag': 'tag2', 'name': 'name2', 'value': 6.0, 'timestamp': '2023-01-03'},
            {'tag': 'tag1', 'name': 'name1',
                'value': None, 'timestamp': '2023-01-04'},
        ],
        'model_tags': {
            'name1': {'aggr_function': 'avg'},
            'name2': {'aggr_function': 'max'},
        }
    }

    try:
        await gates_fixture.aggregate_data(input_data)
    except Exception as e:
        assert str(e) == "Test exception"
        gates_fixture.notification_handler.build_and_send_notification.assert_called_once_with(
            notification_id="AGGREGATION_ISSUES",
            message="Error aggregating data: Test exception",
            block="aggregate_data",
            level=NotificationLevel.ERROR,
            attachment_content=ANY
        )
    else:
        assert False