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() gates = Gates(logger=logger, notification_handler=notification_handler) gates.send_notification = MagicMock() return gates metadata = { 'metadata': { 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', 'workflow_name': 'scouter' } } @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': { 'policy': 'DISCARD' } }, 'data': { 'name': ['tag1', 'tag2', 'tag3'], '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]} }, **metadata } # Execute result = await gates_fixture.data_quality_gate(input_data) # Verify assert len(result['tag']) == 2 assert 'tag2' not in result['tag'] gates_fixture.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': { 'policy': 'KEEP' } }, 'data': { 'name': ['tag1', 'tag2', 'tag3'], '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]} }, **metadata } # 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.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': { 'policy': 'DISCARD' }, 'OUT_OF_BOUNDS_FILTER': { 'policy': '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]} }, **metadata } 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.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 gates_fixture.warning = MagicMock() input_data = { 'filters': { 'UNKNOWN_FILTER': { 'policy': 'DISCARD' } }, 'data': { 'tag': ['tag1'], 'name': ['tag1'], 'value': [1.0], 'timestamp': ['2023-01-01'] }, 'model_tags': { 'tag1': {'data_range': [0, 100]} }, **metadata } # 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.warning.assert_called_once_with( "Filter UNKNOWN_FILTER not found", metadata=metadata['metadata'] ) @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': { 'policy': 'DISCARD' } }, 'data': { 'tag': ['tag1'], 'name': ['tag1'], 'value': [1.0], 'timestamp': ['2023-01-01'] }, 'model_tags': { 'tag1': {'data_range': [0, 100]} }, **metadata } # 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.send_notification.assert_called_once() call_args = gates_fixture.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': { 'policy': 'DISCARD' } }, 'data': { 'tag': [], 'name': [], 'value': [], 'timestamp': [] }, 'model_tags': {}, **metadata } # Execute result = await gates_fixture.data_quality_gate(input_data) # Verify empty result and no notifications assert len(result['tag']) == 0 gates_fixture.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]} }, **metadata } # Execute result = await gates_fixture.data_quality_gate(input_data) # Verify data is unchanged and no notifications assert len(result['tag']) == 1 gates_fixture.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, metadata) assert result == expected_result # Check notification was sent for invalid function if aggr_function == 'invalid': gates_fixture.send_notification.assert_called_once() else: gates_fixture.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_func': 'avg'}, 'name2': {'aggr_func': 'max'}, }, **metadata } # 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.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'}, }, **metadata } # Expected result expected_result = {} # Execute result = await gates_fixture.aggregate_data(input_data) # Verify assert result == expected_result gates_fixture.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'}, }, **metadata } try: await gates_fixture.aggregate_data(input_data) except Exception as e: assert str(e) == "Test exception" gates_fixture.send_notification.assert_called_once_with( metadata=metadata['metadata'], notification_id="AGGREGATION_ISSUES", message="Error aggregating data: Test exception", block="aggregate_data", level=NotificationLevel.ERROR, attachment_content=ANY ) else: assert False @pytest.mark.asyncio @patch('scouter.activities.gates.metrics') async def test_write_metrics(mock_metrics, gates_fixture): """Test write_metrics method.""" input_data = { 'metadata': metadata['metadata'] } await gates_fixture.write_metrics(input_data) mock_metrics.LABORIOUS_DATA_WRITTEN_COUNT.labels.assert_called_once_with( pod_id=gates_fixture.pod_id, model_name=metadata['metadata']['model_name'], pipeline_name=metadata['metadata']['workflow_name'] )