diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index 077a433..c279756 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -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) \ No newline at end of file diff --git a/laborious/activities/api.py b/laborious/activities/api.py new file mode 100644 index 0000000..df186d5 --- /dev/null +++ b/laborious/activities/api.py @@ -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() \ No newline at end of file diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index ba8cd4a..4476c09 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -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. diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 33384b0..db85a78 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -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 diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index 6c5f770..40bef4e 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -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 diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index e50c2cf..1aca3c3 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -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( @@ -151,21 +154,37 @@ 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 diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 73f0227..88e9e3c 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -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'], }, ) diff --git a/requirements.txt b/requirements.txt index 3921162..87d3674 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 @@ -12,4 +13,5 @@ s3fs pyarrow kaleido hyperopt -shap \ No newline at end of file +shap +pycurl \ No newline at end of file diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index e0c3ec9..5c979e9 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -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() diff --git a/tests/laborious/activities/test_api.py b/tests/laborious/activities/test_api.py new file mode 100644 index 0000000..38339ea --- /dev/null +++ b/tests/laborious/activities/test_api.py @@ -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() diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 80516ad..80eb19f 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -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' diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index f5619f3..ffac5c6 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -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 diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 6581033..a31ab2a 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -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, diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 6513da5..c19d6ef 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -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(