Files
sientia-dataops-scouter_tem…/tests/activities/test_gates.py
vitor-aignosi 97eb5bc904 SIENTIAPDE-1316
Update .gitignore and refactor metrics.py, activities.py, and gates.py for improved clarity and consistency. Added coverage.xml and cache directories to .gitignore. Standardized string formatting and parameter handling in metrics and activities classes, enhancing code readability. Removed the deprecated faker.py file and adjusted related tests accordingly.
2025-10-16 13:31:12 -03:00

369 lines
13 KiB
Python

from typing import Any
from unittest.mock import ANY, MagicMock, Mock, patch
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: dict[str, Any] = {}
# 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:
raise AssertionError('Exception not raised')
@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'],
)