vitor-aignosi
1bddde17f4
Refactor Activities and API Integration for PI Web API - Reintroduced the API import in the Activities class for proper integration. - Cleaned up whitespace and formatting in the API class and related tests for improved readability. - Updated test cases to ensure consistent formatting in error messages and configuration structures for PI Web API. - Enhanced connectors_config.py with additional whitespace for better organization.
254 lines
7.7 KiB
Python
254 lines
7.7 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_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()
|
|
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()
|
|
|
|
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()
|
|
|
|
api.pi_web_api_client.write_value.side_effect = [
|
|
None,
|
|
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()
|
|
|
|
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()
|