Code import - branch 0.5.0
This commit is contained in:
492
tests/laborious/activities/test_api.py
Normal file
492
tests/laborious/activities/test_api.py
Normal file
@@ -0,0 +1,492 @@
|
||||
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']
|
||||
)
|
||||
Reference in New Issue
Block a user