vitor-aignosi
00b6382a88
Simplify bandit nosec comments and mock metrics_controller with real spe - drift.py/simple_metrics.py: move nosec rationale to a code comment above the query instead of an inline suppression comment, keeping bandit happy with a plain `# nosec` - tests: replace bare `MagicMock()` with `MagicMock(spec=MetricsController)` for metrics_controller across activities/repository tests so mocks fail fast on attribute typos instead of silently accepting any call
542 lines
17 KiB
Python
542 lines
17 KiB
Python
from unittest.mock import ANY, MagicMock, call, patch
|
|
|
|
from pytest import fixture
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
|
|
from laborious.activities.api import API, PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'schema_name': 'test_schedule',
|
|
},
|
|
}
|
|
|
|
|
|
def _create_mock_dataframe(to_dict_return=None):
|
|
"""Helper function to create a mocked DataFrame for testing."""
|
|
mock_df = MagicMock()
|
|
mock_head = MagicMock()
|
|
|
|
def get_column_values(key):
|
|
if key == 'prediction':
|
|
return MagicMock(values=[0.75])
|
|
elif key == 'prediction_confidence':
|
|
return MagicMock(values=[0.95])
|
|
else:
|
|
return MagicMock(values=['2024-01-01T00:00:00+00:00'])
|
|
|
|
mock_head.__getitem__.side_effect = get_column_values
|
|
mock_df.head.return_value = mock_head
|
|
|
|
if to_dict_return is None:
|
|
to_dict_return = {
|
|
'prediction': [0.75],
|
|
'prediction_confidence': [0.95],
|
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
|
}
|
|
mock_df.to_dict.return_value = to_dict_return
|
|
|
|
return mock_df
|
|
|
|
|
|
@fixture
|
|
def base_input_data():
|
|
"""Base input data for PI Web API tests."""
|
|
return {
|
|
**metadata,
|
|
'data': {
|
|
'prediction': [0.75],
|
|
'prediction_confidence': [0.95],
|
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
|
},
|
|
'pi_web_api_output_config': {
|
|
'endpoint': 'https://test-pi-server.com/piwebapi',
|
|
'prediction_tags': {'tag1': 'web_id_1'},
|
|
'confidence_tags': {'tag2': 'web_id_2'},
|
|
},
|
|
}
|
|
|
|
|
|
@patch('laborious.activities.api.PIWebAPIClient')
|
|
def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_client):
|
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
|
|
|
api_instance = API(
|
|
base_url='https://test-pi-server.com',
|
|
auth_type='bearer',
|
|
auth_token='test_token',
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=MagicMock(spec=MetricsController),
|
|
)
|
|
with patch.object(
|
|
SientiaMonitoring,
|
|
'get_core_labels',
|
|
return_value={
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'local',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'operation_type': 'write_pi_web_api_data',
|
|
},
|
|
):
|
|
labels = api_instance.get_pi_web_api_core_labels(metadata=metadata['metadata'])
|
|
assert labels['operation_type'] == 'write_pi_web_api_data'
|
|
assert labels == {
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'local',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'operation_type': 'write_pi_web_api_data',
|
|
}
|
|
|
|
|
|
@patch('laborious.activities.api.PIWebAPIClient')
|
|
def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client):
|
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
|
|
|
api_instance = API(
|
|
base_url='https://test-pi-server.com',
|
|
auth_type='bearer',
|
|
auth_token='test_token',
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=MagicMock(spec=MetricsController),
|
|
)
|
|
with patch.object(
|
|
SientiaMonitoring,
|
|
'get_core_labels',
|
|
return_value={
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'k8s',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'operation_type': 'write',
|
|
},
|
|
):
|
|
labels = api_instance.get_pi_web_api_core_labels(
|
|
metadata=metadata['metadata'], operation_type='write'
|
|
)
|
|
assert labels['operation_type'] == 'write'
|
|
assert labels['runtime'] == 'k8s'
|
|
|
|
|
|
def test__init__():
|
|
api = API(
|
|
base_url='https://test-pi-server.com',
|
|
auth_type='bearer',
|
|
auth_token='test_token',
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=MagicMock(spec=MetricsController),
|
|
)
|
|
|
|
assert api.pi_web_api_client is not None
|
|
|
|
|
|
@fixture
|
|
@patch('laborious.activities.api.PIWebAPIClient')
|
|
def api(mock_pi_web_api_client):
|
|
mock_client = MagicMock()
|
|
mock_client.write_value = MagicMock()
|
|
mock_client.close = MagicMock()
|
|
mock_client.base_url = 'https://test-pi-server.com'
|
|
mock_pi_web_api_client.return_value = mock_client
|
|
|
|
api_instance = API(
|
|
base_url='https://test-pi-server.com',
|
|
auth_type='bearer',
|
|
auth_token='test_token',
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=MagicMock(spec=MetricsController),
|
|
)
|
|
api_instance.send_notification = MagicMock()
|
|
api_instance.info = MagicMock()
|
|
api_instance.error = MagicMock()
|
|
api_instance.emit_metric_sync = MagicMock()
|
|
api_instance.get_core_labels = MagicMock(
|
|
return_value={
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'local',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
)
|
|
return api_instance
|
|
|
|
|
|
@patch('laborious.activities.api.DataFrame')
|
|
def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_data):
|
|
input_data = {
|
|
**base_input_data,
|
|
'pi_web_api_output_config': {
|
|
'endpoint': 'https://test-pi-server.com/piwebapi',
|
|
'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
|
|
'confidence_tags': {'tag3': 'web_id_3', 'tag4': 'web_id_4'},
|
|
},
|
|
}
|
|
|
|
mock_dataframe.return_value = _create_mock_dataframe()
|
|
|
|
# Mock successful responses
|
|
api.pi_web_api_client.write_value.side_effect = [
|
|
[{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': []}],
|
|
[{'WebId': 'web_id_3', 'Errors': []}, {'WebId': 'web_id_4', 'Errors': []}],
|
|
]
|
|
|
|
result = api.write_pi_web_api_data(input_data)
|
|
|
|
api.pi_web_api_client.write_value.assert_has_calls(
|
|
[
|
|
call(
|
|
web_ids=['web_id_1', 'web_id_2'],
|
|
value={
|
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
|
'Value': 0.75,
|
|
},
|
|
metadata=metadata['metadata'],
|
|
),
|
|
call(
|
|
web_ids=['web_id_3', 'web_id_4'],
|
|
value={
|
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
|
'Value': 0.95,
|
|
},
|
|
metadata=metadata['metadata'],
|
|
),
|
|
]
|
|
)
|
|
|
|
assert result == {
|
|
'prediction': [0.75],
|
|
'prediction_confidence': [0.95],
|
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
|
}
|
|
|
|
|
|
@patch('laborious.activities.api.DataFrame')
|
|
def test_write_pi_web_api_data_partial_prediction_mismatch_downgrades_confidence(
|
|
mock_dataframe, api, base_input_data
|
|
):
|
|
"""A partial write (some tags report Errors) doesn't raise — it's caught by the tag-count
|
|
mismatch check, which still must downgrade confidence and attach the mismatch message."""
|
|
input_data = {
|
|
**base_input_data,
|
|
'pi_web_api_output_config': {
|
|
'endpoint': 'https://test-pi-server.com/piwebapi',
|
|
'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
|
|
'confidence_tags': {'tag3': 'web_id_3'},
|
|
},
|
|
}
|
|
mock_df = _create_mock_dataframe()
|
|
mock_dataframe.return_value = mock_df
|
|
|
|
api.pi_web_api_client.write_value.side_effect = [
|
|
[{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': ['bad tag']}],
|
|
[{'WebId': 'web_id_3', 'Errors': []}],
|
|
]
|
|
|
|
api.write_pi_web_api_data(input_data)
|
|
|
|
mock_df.__setitem__.assert_any_call(
|
|
'prediction_confidence', PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
|
)
|
|
mock_df.__setitem__.assert_any_call('comments', ANY)
|
|
api.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
|
message=ANY,
|
|
block='write_pi_web_api_data',
|
|
level=NotificationLevel.ERROR,
|
|
)
|
|
|
|
|
|
@patch('laborious.activities.api.DataFrame')
|
|
def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
|
|
mock_dataframe.return_value = _create_mock_dataframe(
|
|
{
|
|
'prediction': [0.75],
|
|
'prediction_confidence': [PI_WEB_API_PREDICTION_ERROR_CONFIDENCE],
|
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
|
}
|
|
)
|
|
|
|
api.pi_web_api_client.write_value.side_effect = Exception('Prediction write failed')
|
|
|
|
result = api.write_pi_web_api_data(base_input_data)
|
|
|
|
api.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
|
message="Error writing prediction data to PI Web API: Prediction write failed\n Tags: {'tag1': 'web_id_1'}",
|
|
block='write_pi_web_api_data',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
assert result['prediction_confidence'][0] == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
|
assert api.pi_web_api_client.write_value.call_count == 1
|
|
|
|
|
|
@patch('laborious.activities.api.DataFrame')
|
|
def test_write_pi_web_api_data_confidence_error(mock_dataframe, api, base_input_data):
|
|
mock_dataframe.return_value = _create_mock_dataframe()
|
|
|
|
# First call succeeds, second fails
|
|
api.pi_web_api_client.write_value.side_effect = [
|
|
[{'WebId': 'web_id_1', 'Errors': []}],
|
|
Exception('Confidence write failed'),
|
|
]
|
|
|
|
result = api.write_pi_web_api_data(base_input_data)
|
|
|
|
api.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
|
|
message="Error writing confidence data to PI Web API: Confidence write failed\n Tags: {'tag2': 'web_id_2'}",
|
|
block='write_pi_web_api_data',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
assert result == {
|
|
'prediction': [0.75],
|
|
'prediction_confidence': [0.95],
|
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
|
}
|
|
assert api.pi_web_api_client.write_value.call_count == 2
|
|
|
|
|
|
@patch('laborious.activities.api.DataFrame')
|
|
def test_write_pi_web_api_data_empty_tags(mock_dataframe, api, base_input_data):
|
|
input_data = {
|
|
**base_input_data,
|
|
'pi_web_api_output_config': {
|
|
'endpoint': 'https://test-pi-server.com/piwebapi',
|
|
'prediction_tags': {},
|
|
'confidence_tags': {},
|
|
},
|
|
}
|
|
|
|
mock_dataframe.return_value = _create_mock_dataframe()
|
|
|
|
# Mock empty responses
|
|
api.pi_web_api_client.write_value.side_effect = [
|
|
[],
|
|
[],
|
|
]
|
|
|
|
result = api.write_pi_web_api_data(input_data)
|
|
|
|
api.pi_web_api_client.write_value.assert_has_calls(
|
|
[
|
|
call(
|
|
web_ids=[],
|
|
value={
|
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
|
'Value': 0.75,
|
|
},
|
|
metadata=metadata['metadata'],
|
|
),
|
|
call(
|
|
web_ids=[],
|
|
value={
|
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
|
'Value': 0.95,
|
|
},
|
|
metadata=metadata['metadata'],
|
|
),
|
|
]
|
|
)
|
|
|
|
assert result == {
|
|
'prediction': [0.75],
|
|
'prediction_confidence': [0.95],
|
|
'timestamp': ['2024-01-01T00:00:00+00:00'],
|
|
}
|
|
|
|
|
|
@patch('laborious.activities.api.DataFrame')
|
|
def test_write_pi_web_api_data_updates_confidence_and_comments(
|
|
mock_dataframe, api, base_input_data
|
|
):
|
|
mock_dataframe.return_value = _create_mock_dataframe()
|
|
api.pi_web_api_client.write_value.side_effect = [
|
|
[{'WebId': 'web_id_1', 'Errors': []}],
|
|
[{'WebId': 'web_id_2', 'Errors': []}],
|
|
]
|
|
with patch.object(
|
|
api,
|
|
'process_pi_web_api_response',
|
|
new=MagicMock(side_effect=[(0.33, 'PI warning'), (0, '')]),
|
|
) as process_mock:
|
|
result = api.write_pi_web_api_data(base_input_data)
|
|
|
|
assert process_mock.call_count == 2
|
|
assert result is not None
|
|
|
|
|
|
@patch('laborious.activities.api.SientiaMonitoring.shutdown')
|
|
def test_close(mock_shutdown, api):
|
|
api.close()
|
|
|
|
api.pi_web_api_client.close.assert_called_once()
|
|
mock_shutdown.assert_called_once_with(api)
|
|
|
|
|
|
def test_process_pi_web_api_response_success(api):
|
|
"""Test successful processing of PI Web API response with all tags written."""
|
|
response_data = [
|
|
{'WebId': 'web_id_1', 'Errors': []},
|
|
{'WebId': 'web_id_2', 'Errors': []},
|
|
]
|
|
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
|
|
core_labels = {
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'local',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence, message = api.process_pi_web_api_response(
|
|
response_data=response_data,
|
|
tags=tags,
|
|
core_labels=core_labels,
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert confidence == 0
|
|
assert message == ''
|
|
assert api.emit_metric_sync.call_count == 2
|
|
# Verify that emit_metric_sync was called with correct tags structure
|
|
call_args_list = api.emit_metric_sync.call_args_list
|
|
assert len(call_args_list) == 2
|
|
# Check that all calls include core_labels and tag_name
|
|
for call_args in call_args_list:
|
|
assert 'tag_name' in call_args.kwargs['tags']
|
|
assert call_args.kwargs['tags']['tag_name'] in ['tag1', 'tag2']
|
|
|
|
|
|
def test_process_pi_web_api_response_with_errors(api):
|
|
"""Test processing response with errors in some tags."""
|
|
response_data = [
|
|
{'WebId': 'web_id_1', 'Errors': ['Error writing tag']},
|
|
{'WebId': 'web_id_2', 'Errors': []},
|
|
]
|
|
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
|
|
core_labels = {
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'local',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence, message = api.process_pi_web_api_response(
|
|
response_data=response_data,
|
|
tags=tags,
|
|
core_labels=core_labels,
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
|
assert (
|
|
message
|
|
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
|
|
)
|
|
assert api.emit_metric_sync.call_count == 2
|
|
|
|
|
|
def test_process_pi_web_api_response_missing_tags(api):
|
|
"""Test processing response when number of written tags doesn't match expected."""
|
|
response_data = [
|
|
{'WebId': 'web_id_1', 'Errors': []},
|
|
]
|
|
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
|
|
core_labels = {
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'local',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence, message = api.process_pi_web_api_response(
|
|
response_data=response_data,
|
|
tags=tags,
|
|
core_labels=core_labels,
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
|
assert (
|
|
message
|
|
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag1'] tags were written."
|
|
)
|
|
api.send_notification.assert_called_once()
|
|
call_args = api.send_notification.call_args
|
|
assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
|
|
assert call_args.kwargs['level'] == NotificationLevel.ERROR
|
|
|
|
|
|
def test_process_pi_web_api_response_missing_webid(api):
|
|
"""Test processing response when WebId is missing in response item."""
|
|
response_data = [
|
|
{'Errors': []},
|
|
{'WebId': 'web_id_2', 'Errors': []},
|
|
]
|
|
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
|
|
core_labels = {
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'local',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence, message = api.process_pi_web_api_response(
|
|
response_data=response_data,
|
|
tags=tags,
|
|
core_labels=core_labels,
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
|
assert (
|
|
message
|
|
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
|
|
)
|
|
api.error.assert_any_call('The response did not contain some WebIds', metadata['metadata'])
|
|
|
|
|
|
def test_process_pi_web_api_response_missing_tag_name(api):
|
|
"""Test processing response when tag name is not found for WebId."""
|
|
response_data = [
|
|
{'WebId': 'unknown_web_id', 'Errors': []},
|
|
]
|
|
tags = {'tag1': 'web_id_1'}
|
|
core_labels = {
|
|
'pod_id': 'test_pod',
|
|
'runtime': 'local',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence, message = api.process_pi_web_api_response(
|
|
response_data=response_data,
|
|
tags=tags,
|
|
core_labels=core_labels,
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
|
assert (
|
|
message
|
|
== "The number of written tags does not match the number of tag names: Expected ['tag1'] tags, but [] tags were written."
|
|
)
|
|
api.error.assert_any_call(
|
|
'The response did not contain the tag name for WebId unknown_web_id', metadata['metadata']
|
|
)
|