vitor-aignosi
241f283724
Enhance end-to-end tests for PredictionsBatch workflow scenarios - Introduced mock repositories for PI Web API and OPC operations to improve test coverage. - Updated test scenarios to handle partial write errors for PI Web API and OPC. - Refactored existing tests to assert correct behavior under various error conditions. - Enhanced logging and error handling in API and OPC activities to provide clearer feedback on failures. - Removed outdated integration test file to streamline test suite.
426 lines
14 KiB
Python
426 lines
14 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.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, 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 = {
|
|
'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, 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 = {
|
|
'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, 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 = {
|
|
'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, 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 = {
|
|
'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, 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']
|
|
)
|