vitor-aignosi
ecbcd67597
Refactor PI Web API response handling in tests - Updated test cases in test_api.py to handle response data as lists instead of dictionaries for consistency with the API's expected output format. - Adjusted mock responses to reflect the new structure, ensuring tests accurately simulate API behavior. - Enhanced clarity in test descriptions and improved overall test coverage for response processing scenarios.
423 lines
13 KiB
Python
423 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 = [
|
|
[{'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',
|
|
'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',
|
|
'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',
|
|
'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',
|
|
'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',
|
|
'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']
|
|
)
|