vitor-aignosi
e17824eb85
Refactor metrics and API handling for improved consistency and clarity - Removed the SIENTIA_CORE_LABELS constant and replaced it with CORE_LABELS for uniformity across metrics. - Updated the API class to ensure operation_type is always included in core labels for PI Web API metrics. - Simplified metric tag handling in the Gates class by consolidating common tags into a single core_tags dictionary. - Enhanced the PredictionProcess class to improve error handling and variable naming for clarity.
493 lines
15 KiB
Python
493 lines
15 KiB
Python
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
|
|
|
import pytest_asyncio
|
|
from pytest import fixture, mark
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
|
|
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=AsyncMock(),
|
|
)
|
|
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=AsyncMock(),
|
|
)
|
|
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=AsyncMock(),
|
|
)
|
|
|
|
assert api.pi_web_api_client is not None
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
@patch('laborious.activities.api.PIWebAPIClient')
|
|
def api(mock_pi_web_api_client):
|
|
mock_client = MagicMock()
|
|
mock_client.write_value = AsyncMock()
|
|
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=AsyncMock(),
|
|
)
|
|
api_instance.send_notification_async = AsyncMock()
|
|
api_instance.info = MagicMock()
|
|
api_instance.error = MagicMock()
|
|
api_instance.emit_metric = AsyncMock()
|
|
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
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.api.DataFrame')
|
|
async 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 = await 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'],
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.api.DataFrame')
|
|
async 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 = await api.write_pi_web_api_data(base_input_data)
|
|
|
|
api.send_notification_async.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
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.api.DataFrame')
|
|
async 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 = await api.write_pi_web_api_data(base_input_data)
|
|
|
|
api.send_notification_async.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
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.api.DataFrame')
|
|
async 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 = await 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'],
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_close(api):
|
|
api.close()
|
|
|
|
api.pi_web_api_client.close.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async 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 = await 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.call_count == 2
|
|
# Verify that emit_metric was called with correct tags structure
|
|
call_args_list = api.emit_metric.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']
|
|
|
|
|
|
@mark.asyncio
|
|
async 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 = await 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.call_count == 2
|
|
|
|
|
|
@mark.asyncio
|
|
async 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 = await 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_async.assert_called_once()
|
|
call_args = api.send_notification_async.call_args
|
|
assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
|
|
assert call_args.kwargs['level'] == NotificationLevel.ERROR
|
|
|
|
|
|
@mark.asyncio
|
|
async 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 = await 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'])
|
|
|
|
|
|
@mark.asyncio
|
|
async 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 = await 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']
|
|
)
|