Code import - branch feature/SIENTIAPDE-1646
This commit is contained in:
409
tests/activities/test_gates.py
Normal file
409
tests/activities/test_gates.py
Normal file
@@ -0,0 +1,409 @@
|
||||
from typing import Any
|
||||
from unittest.mock import ANY, MagicMock, Mock, call, 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()
|
||||
metrics_controller = MagicMock()
|
||||
gates = Gates(
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
gates.send_notification = MagicMock()
|
||||
|
||||
gates.logger = logger
|
||||
gates.notification_handler = notification_handler
|
||||
gates.metrics_controller = metrics_controller
|
||||
gates.pod_id = 'localhost'
|
||||
return gates
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'scouter',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@patch('scouter.activities.gates.SientiaMonitoring')
|
||||
def test_close(mock_sientia_monitoring, gates_fixture):
|
||||
"""Test close method."""
|
||||
gates_fixture.close()
|
||||
mock_sientia_monitoring.shutdown.assert_called_once()
|
||||
|
||||
|
||||
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 = 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()
|
||||
|
||||
|
||||
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 = 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()
|
||||
|
||||
|
||||
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 = 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
|
||||
|
||||
|
||||
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 = 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']
|
||||
)
|
||||
|
||||
|
||||
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 = 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']
|
||||
|
||||
|
||||
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 = 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()
|
||||
|
||||
|
||||
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 = 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'),
|
||||
(pd.DataFrame({'value': [10.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()
|
||||
|
||||
|
||||
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 = gates_fixture.aggregate_data(input_data)
|
||||
|
||||
# Verify
|
||||
assert result == expected_result
|
||||
gates_fixture.send_notification.assert_not_called()
|
||||
|
||||
|
||||
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 = gates_fixture.aggregate_data(input_data)
|
||||
|
||||
# Verify
|
||||
assert result == expected_result
|
||||
gates_fixture.send_notification.assert_not_called()
|
||||
|
||||
|
||||
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:
|
||||
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')
|
||||
|
||||
|
||||
@patch('scouter.activities.gates.metrics')
|
||||
def test_write_metrics(mock_metrics, gates_fixture):
|
||||
"""Test write_metrics method."""
|
||||
input_data = {
|
||||
'metadata': metadata['metadata'],
|
||||
'tag_values': {
|
||||
'variable': ['tag1', 'tag2', 'tag3'],
|
||||
'value': [1.0, 2.0, None],
|
||||
},
|
||||
}
|
||||
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'],
|
||||
workflow_name=metadata['metadata']['workflow_name'],
|
||||
)
|
||||
mock_metrics.LABORIOUS_DATA_WRITTEN_COUNT.labels.return_value.inc.assert_called_once()
|
||||
|
||||
# Only tags with None values should be registered
|
||||
mock_metrics.TAG_CHANGES_MONITOR.labels.return_value.set.assert_has_calls(
|
||||
[
|
||||
call(1.0),
|
||||
call(2.0),
|
||||
],
|
||||
any_order=True,
|
||||
)
|
||||
|
||||
mock_metrics.TAG_CHANGES_MONITOR.labels.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
pod_id=gates_fixture.pod_id,
|
||||
model_name=metadata['metadata']['model_name'],
|
||||
workflow_name=metadata['metadata']['workflow_name'],
|
||||
tag_name='tag1',
|
||||
),
|
||||
call(
|
||||
pod_id=gates_fixture.pod_id,
|
||||
model_name=metadata['metadata']['model_name'],
|
||||
workflow_name=metadata['metadata']['workflow_name'],
|
||||
tag_name='tag2',
|
||||
),
|
||||
],
|
||||
any_order=True,
|
||||
)
|
||||
Reference in New Issue
Block a user