gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1478

Browse Source 
Enhance Activities and Prediction Workflows with PI Web API Integration

- Updated the Activities class to include API integration, allowing for configuration of PI Web API parameters.
- Modified prediction workflows to support output configuration for PI Web API, enabling data writing to the API.
- Refactored connectors_config.py by removing unused PostgreSQL and MongoDB configuration functions.
- Added tests to validate the new PI Web API functionality in activities and workflows, ensuring robust integration and functionality.
This commit is contained in:
vitor-aignosi
2026-01-08 16:04:19 -03:00
parent 33874ea1fb
commit 892823df11
14 changed files with 783 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
laborious/activities/activities.py
View File

@@ -12,9 +12,10 @@ with workflow.unsafe.imports_passed_through():
from laborious.activities.model_metrics import ModelMetrics
from laborious.activities.opc import OPC
from laborious.activities.storage import Storage
from laborious.activities.api import API
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
"""
Main activities orchestrator for the Laborious system.
@@ -42,6 +43,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
mlflow_config: dict[str, Any],
minio_config: dict[str, Any],
opc_config: dict[str, Any],
pi_web_api_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler,
):
@@ -58,6 +60,8 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
Required keys: host, port, username, password
opc_config: OPC server configuration dictionary
Can contain multiple server configurations
pi_web_api_config: PI Web API server configuration dictionary
Required keys: base_url, auth_type, auth_token
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
@@ -116,6 +120,16 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
metrics_controller=metrics_controller,
)
API.__init__(
self,
base_url=pi_web_api_config['base_url'],
auth_type=pi_web_api_config['auth_type'],
auth_token=pi_web_api_config['auth_token'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
async def shutdown(self):
"""
Gracefully shutdown all activities and clean up resources.
@@ -133,3 +147,4 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
Gates.close(self)
await OPC.close(self)
ModelMetrics.close(self)
API.close(self)

142
laborious/activities/api.py Normal file
View File

@@ -0,0 +1,142 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from typing import Any
from pandas import DataFrame
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.repository.pi_web_api_client import PIWebAPIClient
PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13
class API(SientiaMonitoring):
"""
PI Web API operations for writing data to PI Web API.
This class provides Temporal activities for interacting with the PI Web API
to write data to PI Web API.
"""
def __init__(
self,
base_url: str,
auth_type: str,
auth_token: str,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
) -> None:
"""
Initialize API activity with PI Web API client.
Args:
base_url (str): Base URL of the PI Web API server
auth_type (str): Authentication type ('basic' or 'bearer')
auth_token (str): Authentication token
logger (Logger): Logger instance for operation logging
notification_handler (NotificationHandler): Handler for system notifications
metrics_controller (MetricsController): Controller for metrics collection
"""
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.pi_web_api_client = PIWebAPIClient(
base_url=base_url,
auth_config={
'type': auth_type,
'token': auth_token,
},
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
def close(self) -> None:
"""
Close the PI Web API client and shutdown monitoring services.
"""
self.pi_web_api_client.close()
SientiaMonitoring.shutdown(self)
@activity.defn(name='write_pi_web_api_data')
async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Write data to PI Web API.
Args:
input_data (dict[str, Any]): The input data. Containing:
- metadata (dict[str, Any]): The metadata.
- pi_web_api_output_config (dict[str, Any]): The PI Web API output configuration.
- data (dict[str, Any]): The data to write.
"""
metadata = input_data['metadata']
data = DataFrame(input_data['data'])
pi_web_api_output_config = input_data['pi_web_api_output_config']
self.info('Writing data to PI Web API...', metadata)
endpoint = pi_web_api_output_config['endpoint']
raw_prediction_tags = pi_web_api_output_config['prediction_tags']
raw_confidence_tags = pi_web_api_output_config['confidence_tags']
prediction_tags = list[str](raw_prediction_tags.values())
confidence_tags = list[str](raw_confidence_tags.values())
prediction_value = data.head(1)['prediction'].values[0]
confidence_value = data.head(1)['prediction_confidence'].values[0]
try:
await self.pi_web_api_client.write_value(
web_ids=prediction_tags,
value={
'Timestamp': data.head(1)['timestamp'].values[0],
'Value': prediction_value,
},
endpoint=endpoint,
metadata=metadata,
)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}',
block='write_pi_web_api_data',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
return data.to_dict()
try:
await self.pi_web_api_client.write_value(
web_ids=confidence_tags,
value={
'Timestamp': data.head(1)['timestamp'].values[0],
'Value': confidence_value,
},
endpoint=endpoint,
metadata=metadata,
)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}',
block='write_pi_web_api_data',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
return data.to_dict()

64
laborious/utils/connectors_config.py
View File

@@ -2,38 +2,6 @@ import json
from os import getenv
from typing import Any
def build_postgres_config() -> dict[str, Any]:
"""
Build PostgreSQL database configuration from environment variables.
This function constructs a PostgreSQL configuration dictionary from
environment variables with sensible defaults for local development.
It handles connection pool configuration and security parameters.
Environment Variables:
POSTGRES_HOST: Database hostname (default: localhost)
POSTGRES_PORT: Database port (default: 5432)
POSTGRES_USER: Database username (default: sientia)
POSTGRES_PASSWORD: Database password (default: sientia)
POSTGRES_DBNAME: Database name (default: sientia)
POSTGRES_MIN_CONNECTIONS: Minimum connection pool size (default: 5)
POSTGRES_MAX_CONNECTIONS: Maximum connection pool size (default: 20)
Returns:
dict: PostgreSQL configuration dictionary with all required parameters
"""
return {
'host': getenv('POSTGRES_HOST', 'localhost'),
'port': int(getenv('POSTGRES_PORT', '5432')),
'user': getenv('POSTGRES_USER', 'sientia'),
'password': getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')),
}
def build_mlflow_config() -> dict[str, Any]:
"""
Build MLFlow server configuration from environment variables.
@@ -98,38 +66,6 @@ def build_opc_config() -> dict[str, Any]:
}
}
def build_mongodb_config() -> dict[str, Any]:
"""
Build MongoDB configuration from environment variables.
This function constructs a MongoDB configuration dictionary from
environment variables with sensible defaults for local development.
It handles connection string and database name configuration.
Environment Variables:
MONGODB_USERNAME: MongoDB username (default: root)
MONGODB_PASSWORD: MongoDB password (default: wKZDbMNU1c)
MONGODB_URL: MongoDB connection URI (default: localhost:27018)
MONGODB_DATABASE_NAME: MongoDB database name (default: sientia)
MONGODB_TTL_INDEX_HOURS: TTL index duration in hours (default: 1)
Returns:
dict: MongoDB configuration dictionary with connection parameters
"""
username = getenv('MONGODB_USERNAME', 'root')
password = getenv('MONGODB_PASSWORD', 'wKZDbMNU1c')
uri = getenv('MONGODB_URL', 'localhost:27018')
connection_string = f'mongodb://{username}:{password}@{uri}'
return {
'connection_string': connection_string,
'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
}
def build_minio_config() -> dict[str, Any]:
"""
Build MinIO (S3-compatible) configuration from environment variables.

4
laborious/worker/worker.py
View File

@@ -67,9 +67,11 @@ with workflow.unsafe.imports_passed_through():
from laborious.utils.connectors_config import (
build_minio_config,
build_mlflow_config,
build_mongodb_config,
build_opc_config,
)
from sientia_do.connectors_config import (
build_postgres_config,
build_mongodb_config,
)
from laborious.workflows.drift import Drift
from laborious.workflows.minimal_retrain import MinimalRetrain

2
laborious/workflows/predictions_batch.py
View File

@@ -111,7 +111,9 @@ class PredictionsBatch:
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {}),
'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}),
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
'save_transform': input_data.get('save_transform', True),
}
# Execute prediction process workflow

42
laborious/workflows/sub_workflows/format_and_export_prediction.py
View File

@@ -87,6 +87,9 @@ class FormatAndExportPrediction:
transformed_data = input_data.get('transformed_data', None)
prediction_confidence = input_data['prediction_confidence']
opc_output_config = input_data.get('opc_output_config', None)
pi_web_api_output_config = input_data.get('pi_web_api_output_config', None)
if path_flag is None:
# Normal prediction path: format prediction data with full metadata
prediction = await workflow.execute_local_activity_method(
@@ -152,20 +155,36 @@ class FormatAndExportPrediction:
write_transformed_handler = None
opc_metrics = {}
# write to pi web api
if pi_web_api_output_config is not None:
prediction = await workflow.execute_activity_method(
Activities.write_pi_web_api_data,
{
'pi_web_api_output_config': pi_web_api_output_config,
'data': prediction,
**metadata,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
# write to opc
prediction, opc_metrics = await workflow.execute_activity_method(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': prediction,
**metadata,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
if opc_output_config is not None:
prediction, opc_metrics = await workflow.execute_activity_method(
Activities.write_opc_data,
{
'opc_output_config': opc_output_config,
'data': prediction,
**metadata,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
# write to postgres
prediction_handler = workflow.execute_activity_method(
await workflow.execute_activity_method(
Activities.export_data_to_postgres,
{
**metadata,
@@ -178,7 +197,6 @@ class FormatAndExportPrediction:
start_to_close_timeout=timedelta(seconds=180),
)
await prediction_handler
if write_transformed_handler is not None:
await write_transformed_handler

2
laborious/workflows/sub_workflows/prediction_process.py
View File

@@ -207,6 +207,7 @@ class PredictionProcess:
'model_name': model_name,
'model_config': model_config,
'opc_output_config': input_data['opc_output_config'],
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_table_name'],
@@ -294,6 +295,7 @@ class PredictionProcess:
'transform_table_name': transform_table_name,
'comment': comment,
'opc_output_config': input_data['opc_output_config'],
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'prediction_store_policy': input_data['prediction_store_policy'],
},
)

4
requirements.txt
View File

@@ -3,7 +3,8 @@ psycopg2-binary
sqlalchemy
asyncua
redis
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1
#git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1
/home/grezewave/Documents/projects/sientia/sientia-dataops-library/
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.6
prometheus-client
botocore
@@ -13,3 +14,4 @@ pyarrow
kaleido
hyperopt
shap
pycurl

58
tests/laborious/activities/test_activities.py
View File

@@ -3,8 +3,10 @@ from unittest.mock import ANY, AsyncMock, MagicMock, patch
from pytest import mark
from laborious.activities.activities import Activities
from laborious.activities.api import API
from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.model_metrics import ModelMetrics
from laborious.activities.opc import OPC
from laborious.activities.storage import Storage
@@ -13,9 +15,17 @@ from laborious.activities.storage import Storage
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MetricsController')
def test___init__(
mock_metrics_controller, mock_gates_init, mock_opc_init, mock_mlflow_init, mock_storage_init
mock_metrics_controller,
mock_api_init,
mock_model_metrics_init,
mock_gates_init,
mock_opc_init,
mock_mlflow_init,
mock_storage_init,
):
postgres_config = {
'host': 'localhost',
@@ -43,6 +53,12 @@ def test___init__(
'group_id': 'test-group',
}
pi_web_api_config = {
'base_url': 'https://test-pi-server.com',
'auth_type': 'bearer',
'auth_token': 'test_token',
}
logger = MagicMock()
notification_handler = MagicMock()
@@ -51,6 +67,7 @@ def test___init__(
mlflow_config=mlflow_config,
minio_config=minio_config,
opc_config=opc_config,
pi_web_api_config=pi_web_api_config,
logger=logger,
notification_handler=notification_handler,
)
@@ -60,6 +77,8 @@ def test___init__(
assert isinstance(activities, MLFlow)
assert isinstance(activities, OPC)
assert isinstance(activities, Gates)
assert isinstance(activities, ModelMetrics)
assert isinstance(activities, API)
mock_storage_init.assert_called_once_with(
ANY,
@@ -103,13 +122,39 @@ def test___init__(
metrics_controller=mock_metrics_controller.return_value,
)
mock_model_metrics_init.assert_called_once_with(
ANY,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_api_init.assert_called_once_with(
ANY,
base_url=pi_web_api_config['base_url'],
auth_type=pi_web_api_config['auth_type'],
auth_token=pi_web_api_config['auth_token'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
@mark.asyncio
@patch('laborious.activities.activities.Storage')
@patch('laborious.activities.activities.MLFlow')
@patch('laborious.activities.activities.OPC')
@patch('laborious.activities.activities.Gates')
async def test_shutdown(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_storage_init):
@patch('laborious.activities.activities.ModelMetrics')
@patch('laborious.activities.activities.API')
async def test_shutdown(
mock_api_init,
mock_model_metrics_init,
mock_gates_init,
mock_opc_init,
mock_mlflow_init,
mock_storage_init,
):
mock_opc_init.close = AsyncMock()
postgres_config = {
'host': 'localhost',
@@ -137,6 +182,12 @@ async def test_shutdown(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_s
'group_id': 'test-group',
}
pi_web_api_config = {
'base_url': 'https://test-pi-server.com',
'auth_type': 'bearer',
'auth_token': 'test_token',
}
logger = MagicMock()
notification_handler = MagicMock()
@@ -145,6 +196,7 @@ async def test_shutdown(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_s
mlflow_config=mlflow_config,
minio_config=minio_config,
opc_config=opc_config,
pi_web_api_config=pi_web_api_config,
logger=logger,
notification_handler=notification_handler,
)
@@ -154,3 +206,5 @@ async def test_shutdown(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_s
mock_storage_init.close.assert_called_once()
mock_mlflow_init.close.assert_called_once()
mock_gates_init.close.assert_called_once()
mock_model_metrics_init.close.assert_called_once()
mock_api_init.close.assert_called_once()

254
tests/laborious/activities/test_api.py Normal file
View File

@@ -0,0 +1,254 @@
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
import pytest_asyncio
from pandas import DataFrame
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()

75
tests/laborious/utils/test_connectors_config.py
View File

@@ -3,9 +3,7 @@ from os import environ
from laborious.utils.connectors_config import (
build_minio_config,
build_mlflow_config,
build_mongodb_config,
build_opc_config,
build_postgres_config,
)
@@ -92,79 +90,6 @@ def test_build_opc_config_with_defaults():
assert config['1']['reconnection_interval'] == 120
def test_build_postgres_config_with_env_vars():
# Arrange
environ['POSTGRES_HOST'] = 'test-host'
environ['POSTGRES_PORT'] = '5433'
environ['POSTGRES_USER'] = 'test-user'
environ['POSTGRES_PASSWORD'] = 'test-pass'
environ['POSTGRES_DBNAME'] = 'test-db'
environ['POSTGRES_MIN_CONNECTIONS'] = '10'
environ['POSTGRES_MAX_CONNECTIONS'] = '30'
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'test-host'
assert config['port'] == 5433
assert config['user'] == 'test-user'
assert config['password'] == 'test-pass'
assert config['dbname'] == 'test-db'
assert config['min_connections'] == 10
assert config['max_connections'] == 30
def test_build_postgres_config_with_defaults():
# Arrange
environ.pop('POSTGRES_HOST', None)
environ.pop('POSTGRES_PORT', None)
environ.pop('POSTGRES_USER', None)
environ.pop('POSTGRES_PASSWORD', None)
environ.pop('POSTGRES_DBNAME', None)
environ.pop('POSTGRES_MIN_CONNECTIONS', None)
environ.pop('POSTGRES_MAX_CONNECTIONS', None)
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'localhost'
assert config['port'] == 5432
assert config['user'] == 'sientia'
assert config['password'] == 'sientia'
assert config['dbname'] == 'sientia'
assert config['min_connections'] == 5
assert config['max_connections'] == 20
def test_build_mongo_db_config_with_env_vars():
environ['MONGODB_USERNAME'] = 'sientia1'
environ['MONGODB_PASSWORD'] = 'sientia1'
environ['MONGODB_URL'] = 'localhost:27018'
environ['MONGODB_DATABASE_NAME'] = 'test_db'
environ['MONGODB_TTL_INDEX_HOURS'] = '1'
assert build_mongodb_config() == {
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
'database_name': 'test_db',
'ttl_index_seconds': 3600,
}
def test_build_mongo_db_config_with_defaults():
environ.pop('MONGODB_USERNAME', None)
environ.pop('MONGODB_PASSWORD', None)
environ.pop('MONGODB_DATABASE_NAME', None)
environ.pop('MONGODB_URL', None)
environ.pop('MONGODB_TTL_INDEX_HOURS', None)
assert build_mongodb_config() == {
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
'database_name': 'sientia',
'ttl_index_seconds': 3600,
}
def test_build_minio_config_with_env_vars():
environ['MINIO_ENDPOINT_URL'] = 'http://test-host'
environ['MINIO_ACCESS_KEY'] = 'test-key'

266
tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
View File

@@ -375,3 +375,269 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
'path_flag': None,
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
'prediction_store_policy': 'erl:1',
}
pi_web_api_data = MagicMock()
workflow_mock.execute_activity_method.side_effect = [
pi_web_api_data, # write_pi_web_api_data
MagicMock(), # export_data_to_postgres
MagicMock(), # write_metrics
]
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_prediction,
{
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'prediction_store_policy': input_data['prediction_store_policy'],
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_pi_web_api_data,
{
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': pi_web_api_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_metrics,
{
**metadata,
'prediction': pi_web_api_data,
'opc_metrics': {},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_none_path_flag_with_pi_web_api_and_opc(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
'path_flag': None,
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'opc_output_config': {'test': 'config'},
'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
'prediction_store_policy': 'erl:1',
}
prediction_data = MagicMock()
pi_web_api_data = MagicMock()
opc_metrics = MagicMock()
workflow_mock.execute_activity_method.side_effect = [
pi_web_api_data, # write_pi_web_api_data
(prediction_data, opc_metrics), # write_opc_data
MagicMock(), # export_data_to_postgres
MagicMock(), # write_metrics
]
await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_pi_web_api_data,
{
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': pi_web_api_data,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_metrics,
{
**metadata,
'prediction': prediction_data,
'opc_metrics': opc_metrics,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 4
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
'path_flag': 'default',
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
'comment': 'test_comment',
}
pi_web_api_data = MagicMock()
workflow_mock.execute_activity_method.side_effect = [
pi_web_api_data, # write_pi_web_api_data
MagicMock(), # export_data_to_postgres
MagicMock(), # write_metrics
]
await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_pi_web_api_data,
{
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': pi_web_api_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1

5
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

@@ -40,6 +40,7 @@ async def test_run(workflow_mock, prediction_process):
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'},
'pi_web_api_output_config': {'test': 'config'},
'prediction_store_policy': 'lts:1',
}
@@ -183,6 +184,7 @@ async def test_run(workflow_mock, prediction_process):
'model_name': 'test_model_name',
'model_config': input_data['model_config'],
'opc_output_config': input_data['opc_output_config'],
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_table_name'],
@@ -729,6 +731,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'model_name': model_name,
'model_config': model_config,
'opc_output_config': {'test': 'config'},
'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
'prediction_store_policy': prediction_store_policy,
},
confidence,
@@ -755,6 +758,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'transform_table_name': 'test_transform_table',
'comment': 'Prediction Process',
'opc_output_config': {'test': 'config'},
'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
'prediction_store_policy': prediction_store_policy,
},
)
@@ -788,6 +792,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
'model_name': model_name,
'model_config': model_config,
'opc_output_config': {'test': 'config'},
'pi_web_api_output_config': {'endpoint': 'https://test-pi-server.com', 'prediction_tags': {}, 'confidence_tags': {}},
'prediction_store_policy': prediction_store_policy,
},
confidence,

3
tests/laborious/workflows/test_predictions_batch.py
View File

@@ -34,6 +34,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'opc_output_config': 'test_opc_output_config',
'pi_web_api_output_config': 'test_pi_web_api_output_config',
'datetime_columns': ['timestamp', 'created_at'],
'prediction_store_policy': 'erl:1',
'model_config': {'retention': '30'},
@@ -73,7 +74,9 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {}),
'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}),
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
'save_transform': input_data.get('save_transform', True),
}
workflow_mock.execute_child_workflow.assert_has_calls(