vitor-aignosi
7891a7ca6e
SIENTIAPDE-1478 Implement PI Web API response processing and metrics tracking - Added a new method in the API class to process responses from the PI Web API, validating tag writes and emitting metrics for success and errors. - Enhanced error handling for missing WebIds and tag names in responses, with appropriate logging and notifications. - Updated tests to cover various scenarios for processing PI Web API responses, ensuring robust functionality and metrics emission. - Refactored existing methods to integrate the new response processing logic, improving overall code clarity and maintainability.
425 lines
13 KiB
Python
425 lines
13 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'},
|
|
},
|
|
}
|
|
|
|
|
|
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',
|
|
'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 = [
|
|
{'Items': [{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': []}]},
|
|
{'Items': [{'WebId': 'web_id_3', 'Errors': []}, {'WebId': 'web_id_4', 'Errors': []}]},
|
|
]
|
|
|
|
result = await api.write_pi_web_api_data(input_data)
|
|
|
|
api.info.assert_called_once_with('Writing data to PI Web API...', metadata['metadata'])
|
|
|
|
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,
|
|
},
|
|
endpoint='https://test-pi-server.com/piwebapi',
|
|
metadata=metadata['metadata'],
|
|
),
|
|
call(
|
|
web_ids=['web_id_3', 'web_id_4'],
|
|
value={
|
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
|
'Value': 0.95,
|
|
},
|
|
endpoint='https://test-pi-server.com/piwebapi',
|
|
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 = [
|
|
{'Items': [{'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 = [
|
|
{'Items': []},
|
|
{'Items': []},
|
|
]
|
|
|
|
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,
|
|
},
|
|
endpoint='https://test-pi-server.com/piwebapi',
|
|
metadata=metadata['metadata'],
|
|
),
|
|
call(
|
|
web_ids=[],
|
|
value={
|
|
'Timestamp': '2024-01-01T00:00:00+00:00',
|
|
'Value': 0.95,
|
|
},
|
|
endpoint='https://test-pi-server.com/piwebapi',
|
|
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 = {
|
|
'Items': [
|
|
{'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',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence = await api.process_pi_web_api_response(
|
|
response_data=response_data,
|
|
tags=tags,
|
|
core_labels=core_labels,
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert confidence == 0
|
|
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 = {
|
|
'Items': [
|
|
{'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',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence = 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 api.emit_metric.call_count == 2
|
|
api.error.assert_any_call(
|
|
"Error writing tag tag1:web_id_1 to PI Web API: ['Error writing tag']", metadata['metadata']
|
|
)
|
|
|
|
|
|
@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 = {
|
|
'Items': [
|
|
{'WebId': 'web_id_1', 'Errors': []},
|
|
]
|
|
}
|
|
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
|
|
core_labels = {
|
|
'pod_id': 'test_pod',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence = 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
|
|
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 = {
|
|
'Items': [
|
|
{'Errors': []},
|
|
{'WebId': 'web_id_2', 'Errors': []},
|
|
]
|
|
}
|
|
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
|
|
core_labels = {
|
|
'pod_id': 'test_pod',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence = 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
|
|
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 = {
|
|
'Items': [
|
|
{'WebId': 'unknown_web_id', 'Errors': []},
|
|
]
|
|
}
|
|
tags = {'tag1': 'web_id_1'}
|
|
core_labels = {
|
|
'pod_id': 'test_pod',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
}
|
|
|
|
confidence = 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
|
|
api.error.assert_any_call(
|
|
'The response did not contain the tag name for WebId unknown_web_id', metadata['metadata']
|
|
)
|