gitea_admin/sientia-dataops-scouter_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
43512891113d58807e3b54a428ada688efa16c2c
sientia-dataops-scouter_tem…/tests/activities/test_gates.py
vitor-aignosi 4351289111 SIENTIAPDE-1325 
Refactor notification methods in Gates, MongoDB, and Redis classes to use asynchronous send_notification_async. Update related tests to ensure proper async handling and verification of notifications.
2025-11-03 14:57:44 -03:00

422 lines
15 KiB
Python
Raw Blame History

from typing import Any
from unittest.mock import ANY, AsyncMock, 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.send_notification_async = AsyncMock()
gates.emit_metric = AsyncMock()
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()
@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_async.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_async.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_async.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_async.assert_called_once()
call_args = gates_fixture.send_notification_async.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'),
],
)
@pytest.mark.asyncio
async def test_apply_aggregation(gates_fixture, group_data, aggr_function, expected_result):
"""Test apply_aggregation method with various scenarios."""
result = await 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_async.assert_called_once()
else:
gates_fixture.send_notification_async.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 = AsyncMock(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_async.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_async.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'],
'tag_values': {
'variable': ['tag1', 'tag2'],
'value': [1.0, 2.0],
},
}
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'],
)
mock_metrics.LABORIOUS_DATA_WRITTEN_COUNT.labels.return_value.inc.assert_called_once()
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'],
pipeline_name=metadata['metadata']['workflow_name'],
tag_name='tag1',
),
call(
pod_id=gates_fixture.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
tag_name='tag2',
),
],
any_order=True,
)
Reference in New Issue View Git Blame Copy Permalink