From d09fb6ac5ed781e1bc502904d0bb3826d079ec3c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 9 May 2025 16:27:57 -0300 Subject: [PATCH] SIENTIAPDE-994 Refactor activity methods and update requirements.txt to enhance functionality and remove deprecated filters. Added detailed docstrings for clarity and improved error handling in data processing workflows. --- laborious/activities/base.py | 8 + laborious/activities/gates.py | 132 +++++-- laborious/activities/mlflow.py | 24 +- laborious/activities/opc.py | 76 ++-- laborious/activities/postgres.py | 5 +- laborious/utils/filters/base_filter.py | 13 - .../utils/filters/conditional_filters.py | 1 - laborious/utils/filters/mlflow_filters.py | 1 - .../utils/repository/model_repository.py | 2 +- laborious/utils/repository/opc_repository.py | 6 - .../format_and_export_prediction.py | 6 +- .../sub_workflows/prediction_process.py | 106 +++++- requirements.txt | 19 - tests/laborious/activities/test_base.py | 32 ++ tests/laborious/activities/test_gates.py | 335 ++++++++++++++++-- tests/laborious/activities/test_mlflow.py | 121 +++++++ tests/laborious/activities/test_opc.py | 178 ++++++++++ .../repository/test_model_repository.py | 278 +++++++++++++++ .../filters/repository/test_opc_repository.py | 106 ++++++ .../utils/filters/test_conditional_filters.py | 9 +- .../utils/filters/test_mlflow_filters.py | 22 ++ .../test_format_and_export_prediction.py | 115 ++++++ 22 files changed, 1435 insertions(+), 160 deletions(-) delete mode 100644 laborious/utils/filters/base_filter.py create mode 100644 tests/laborious/activities/test_base.py create mode 100644 tests/laborious/activities/test_mlflow.py create mode 100644 tests/laborious/activities/test_opc.py create mode 100644 tests/laborious/utils/filters/repository/test_model_repository.py create mode 100644 tests/laborious/utils/filters/repository/test_opc_repository.py create mode 100644 tests/laborious/utils/filters/test_mlflow_filters.py create mode 100644 tests/laborious/workflows/subworkflows.py/test_format_and_export_prediction.py diff --git a/laborious/activities/base.py b/laborious/activities/base.py index 8c6fba8..9f89e5e 100644 --- a/laborious/activities/base.py +++ b/laborious/activities/base.py @@ -11,6 +11,14 @@ class BaseActivity: def prepare_activity(self, schedule_name: str, model_name: str, model_id: str): + """ + Prepare the activity for the notification handler. + + Args: + schedule_name (str): The name of the schedule. + model_name (str): The name of the model. + model_id (str): The id of the model. + """ self.notification_handler.base_notification.schedule_name = schedule_name self.notification_handler.base_notification.model_name = model_name self.notification_handler.base_notification.model_id = model_id diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index ac07fbf..a2fad98 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -14,15 +14,29 @@ with workflow.unsafe.imports_passed_through(): input_filter_functions = { 'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values, - 'EMPTY_DATA': filter_empty_data + 'EMPTY_DATA': filter_empty_data, + 'path_confidence': { + 'stop': -1, + 'continue': 2, + 'repeat': -1 + } } -transform_filter_functions = { - 'response_filter': { - 'API_ERROR': api_error_filter, +mlflow_response_filter_functions = { + 'API_ERROR': api_error_filter, + 'path_confidence': { + 'stop': -1, + 'continue': 10, + 'repeat': -1 }, - 'content_filter': { - 'NAN_VALUES': nan_values_filter, +} + +mlflow_content_filter_functions = { + 'NAN_VALUES': nan_values_filter, + 'path_confidence': { + 'stop': -1, + 'continue': 18, + 'repeat': -1 } } @@ -40,13 +54,13 @@ class Gates(BaseActivity): input_data (dict): The input data. Contains: filters (dict): The filters to apply. data (dict[str, Any]): The data to filter. + path_priority (list[str]): The path priority. Returns: - tuple[str, int]: ('stop', -1) if some filter policy is 'stop', ('continue', 2) - if no filter policy is 'stop' and some filter policy is 'continue', - None if no filter is applied. + tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions. """ filters = input_data['filters'] data = DataFrame(input_data['data']) + path_priority = input_data['path_priority'] filter_output = [] for fil, config in filters.items(): @@ -63,23 +77,34 @@ class Gates(BaseActivity): attachment_content=trace ) - if 'stop' in filter_output: - return 'stop', -1 - elif 'continue' in filter_output: - return 'continue', 2 + for path_flag in path_priority: + if path_flag in filter_output: + return path_flag, input_filter_functions['path_confidence'][path_flag] return None, 0 - @activity.defn(name="mlflow_gate") - async def mlflow_gate(self, input_data: dict[str, Any]) -> tuple[str, int]: - + @activity.defn(name="mlflow_response_gate") + async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str, int]: + """ + Filters the data based on the mlflow response filters. The return value is a tuple with the first element + being the policy and the second element being the confidence status. + Args: + input_data (dict): The input data. Contains: + filters (dict): The filter configuration to apply. + data (dict[str, Any]): The data to filter. + path_priority (list[str]): The path priority list. + type (str): The type of the gate. + Returns: + tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions. + """ filters = input_data['filters'] - data = DataFrame(input_data['data']) + data = input_data['data'] gate_type = input_data['type'] + path_priority = input_data['path_priority'] filter_output = [] for fil, config in filters.items(): - if transform_filter_functions['response_filter'][fil](data, config): + if mlflow_response_filter_functions[fil](data, config): filter_output.append(config['POLICY']) self.notification_handler.build_and_send_notification( notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}", @@ -89,18 +114,36 @@ class Gates(BaseActivity): attachment_content=data['content']['traceback'] ) - if 'stop' in filter_output: - return 'stop', -1 - elif 'continue' in filter_output: - return 'continue', 10 + for path_flag in path_priority: + if path_flag in filter_output: + return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag] - if gate_type == 'predict': - return None, 0 + return None, 0 - data = DataFrame(data['content']) + @activity.defn(name="mlflow_content_gate") + async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str, int]: + """ + Filters the data based on the mlflow content filters. The return value is a tuple with the first element + being the policy and the second element being the confidence status. + Args: + input_data (dict): The input data. Contains: + filters (dict): The filter configuration to apply. + data (dict[str, Any]): The data to filter. + path_priority (list[str]): The path priority list. + type (str): The type of the gate. + Returns: + tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions. + """ + + filters = input_data['filters'] + data = DataFrame(input_data['data']) + gate_type = input_data['type'] + path_priority = input_data['path_priority'] + + filter_output = [] for fil, config in filters.items(): - if transform_filter_functions['content_filter'][fil](data, config): + if mlflow_content_filter_functions[fil](data, config): filter_output.append(config['POLICY']) self.notification_handler.build_and_send_notification( notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}", @@ -110,17 +153,26 @@ class Gates(BaseActivity): attachment_content=data.to_string() ) - if 'stop' in filter_output: - return 'stop', -1 - elif 'continue' in filter_output: - return 'continue', 18 + for path_flag in path_priority: + if path_flag in filter_output: + return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag] return None, 0 @activity.defn(name="format_prediction") - async def format_prediction(self, input_data: dict[str, Any]) -> str: + async def format_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Formats the prediction data. + Args: + input_data (dict): The input data. Contains: + data (dict[str, Any]): The data to format. + timestamp (str): The timestamp of the data. + model_id (str): The id of the model. + prediction_confidence (float): The confidence of the prediction. + Returns: + dict: The formatted data. + """ data = DataFrame(input_data['data']) - data['timestamp'] = input_data['timestamp'] data['model_id'] = input_data['model_id'] data['prediction_confidence'] = input_data['prediction_confidence'] @@ -131,7 +183,21 @@ class Gates(BaseActivity): return data.to_dict() @activity.defn(name="format_default_prediction") - async def format_default_prediction(self, input_data: dict[str, Any]) -> str: + async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Creates and formats the default prediction data, with zero value in prediction, + and usefull information in the other fields. + + Args: + input_data (dict): The input data. Contains: + timestamp (str): The timestamp of the data. + model_id (str): The id of the model. + prediction_confidence (float): The confidence of the prediction. + comment (str): The comment of the prediction. + Returns: + dict: The formatted data. + """ + return DataFrame({ 'prediction': [0], 'response_time': [0], diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index af3408a..7c785b5 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -8,7 +8,7 @@ with workflow.unsafe.imports_passed_through(): from typing import Any from logging import Logger from sientia_do.notifications.handlers import NotificationHandler - from laborious.utils.model_repository import ModelMonitoringRepository + from laborious.utils.repository.model_repository import MLFlowRepository from sientia_do.notifications.models import NotificationLevel @@ -21,12 +21,22 @@ class MLFlow(BaseActivity): self.mlflow_username = mlflow_username self.mlflow_password = mlflow_password - self.model_monitoring_repository = ModelMonitoringRepository( + self.model_monitoring_repository = MLFlowRepository( f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password ) @activity.defn(name="request_transform") async def request_transform(self, input_data: dict[str, Any]) -> tuple[dict[str, Any], str]: + """ + Access MLFlow model to get the transformed data. + Args: + input_data (dict): The input data. Contains: + data (dict[str, Any]): The data to transform. + model_name (str): The name of the model. + model_retention (int): The retention of the model. + Returns: + tuple[dict[str, Any], str]: The transformed data and the latest timestamp of the data. + """ self.logger.info('Transforming data...') data = DataFrame(input_data['data']) model_name = input_data['model_name'] @@ -50,6 +60,16 @@ class MLFlow(BaseActivity): @activity.defn(name="request_predict") async def request_predict(self, input_data: dict[str, Any]) -> tuple[dict[str, Any], str]: + """ + Access MLFlow model to get the predicted data. + Args: + input_data (dict): The input data. Contains: + data (dict[str, Any]): The data to predict. + model_name (str): The name of the model. + model_retention (int): The retention of the model. + Returns: + tuple[dict[str, Any], str]: The predicted data and the latest timestamp of the data. + """ self.logger.info('Predicting data...') data = DataFrame(input_data['data']) model_name = input_data['model_name'] diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index eae202f..b6d2c38 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -1,5 +1,3 @@ -import traceback -from pandas import DataFrame from temporalio import activity, workflow @@ -10,6 +8,8 @@ with workflow.unsafe.imports_passed_through(): from laborious.utils.repository.opc_repository import OpcRepository from typing import Any from sientia_do.notifications.models import NotificationLevel + import traceback + from pandas import DataFrame class OPC(BaseActivity): @@ -41,36 +41,52 @@ class OPC(BaseActivity): @activity.defn(name='write_opc_data') async def write_opc_data(self, input_data: dict[str, Any]): + """ + Write prediction and confidence data to OPC servers. The two writing + operations are optional and independent of each other. + + Args: + input_data (dict[str, Any]): The input data. Contains the following keys: + data (dict[str, Any]): The dataframe that contains the data to write to the OPC servers. + opc_servers (list[str]): The OPC servers to write to. + opc_output_config (dict[str, Any]): The OPC writing configuration. Contains: + prediction_tags (dict[str, Any]): The tags to write to the OPC servers. + confidence_tags (dict[str, Any]): The tags to write to the OPC servers. + + Returns: + """ data = DataFrame(input_data['data']) _opc_servers = input_data['opc_servers'] opc_output_config = input_data['opc_output_config'] - for tag, config in opc_output_config['prediction_tags'].items(): - try: - self.opc_repository.write_data( - tag, data.head(1)['prediction'].values[0], config['data_type']) - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id="WRITE_OPC_PREDICTION_ERROR", - message=f"Error writing data to OPC server: {e}", - block="write_opc_data", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - self.logger.error(trace) + if 'prediction_tags' in opc_output_config: + for tag, config in opc_output_config['prediction_tags'].items(): + try: + self.opc_repository.write_data( + tag, data.head(1)['prediction'].values[0], config['data_type']) + except Exception as e: + trace = traceback.format_exc() + self.notification_handler.build_and_send_notification( + notification_id="WRITE_OPC_PREDICTION_ERROR", + message=f"Error writing data to OPC server: {e}", + block="write_opc_data", + level=NotificationLevel.ERROR, + attachment_content=trace + ) + self.logger.error(trace) - for tag, config in opc_output_config['confidence_tags'].items(): - try: - self.opc_repository.write_data( - tag, data.head(1)['prediction_confidence'].values[0], config['data_type']) - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id="WRITE_OPC_CONFIDENCE_ERROR", - message=f"Error writing data to OPC server: {e}", - block="write_opc_data", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - self.logger.error(trace) + if 'confidence_tags' in opc_output_config: + for tag, config in opc_output_config['confidence_tags'].items(): + try: + self.opc_repository.write_data( + tag, data.head(1)['prediction_confidence'].values[0], config['data_type']) + except Exception as e: + trace = traceback.format_exc() + self.notification_handler.build_and_send_notification( + notification_id="WRITE_OPC_CONFIDENCE_ERROR", + message=f"Error writing data to OPC server: {e}", + block="write_opc_data", + level=NotificationLevel.ERROR, + attachment_content=trace + ) + self.logger.error(trace) diff --git a/laborious/activities/postgres.py b/laborious/activities/postgres.py index 7ef2309..3716fa9 100644 --- a/laborious/activities/postgres.py +++ b/laborious/activities/postgres.py @@ -135,7 +135,10 @@ class Postgres(BaseActivity): Exports data to a postgres table. Args: - input_data (dict[str, Any]): The data to export. + input_data (dict[str, Any]): The data to export. Contains: + schema (str): The schema of the table. + table_name (str): The name of the table. + data (DataFrame): The data to export. """ schema = input_data["schema"] diff --git a/laborious/utils/filters/base_filter.py b/laborious/utils/filters/base_filter.py deleted file mode 100644 index f5d9948..0000000 --- a/laborious/utils/filters/base_filter.py +++ /dev/null @@ -1,13 +0,0 @@ -from pandas import DataFrame - -class Filter: - def __init__(self, id: str): - self.id = id - self.warnings = [] - - @staticmethod - def method(df: DataFrame) -> DataFrame: - raise NotImplementedError - - def warning(self, message: str): - self.warnings.append(f'[{self.id}] - {message}') \ No newline at end of file diff --git a/laborious/utils/filters/conditional_filters.py b/laborious/utils/filters/conditional_filters.py index d819a3d..6638982 100644 --- a/laborious/utils/filters/conditional_filters.py +++ b/laborious/utils/filters/conditional_filters.py @@ -1,6 +1,5 @@ from typing import List -from laborious.utils.filters.base_filter import Filter from pandas import DataFrame diff --git a/laborious/utils/filters/mlflow_filters.py b/laborious/utils/filters/mlflow_filters.py index 3f5db3b..d054ddb 100644 --- a/laborious/utils/filters/mlflow_filters.py +++ b/laborious/utils/filters/mlflow_filters.py @@ -1,6 +1,5 @@ import numpy as np from pandas import DataFrame -from laborious.utils.filters.base_filter import Filter def api_error_filter(response: dict, _config: dict): diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 5efe83b..3896411 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -16,7 +16,7 @@ from sientia.ModelServing import ModelServing from pathlib import Path -class ModelMonitoringRepository(): +class MLFlowRepository(): def __init__(self, host, username, password): self.model_serving = ModelServing(tracking_uri=host, diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 99fbc08..7d84f85 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/laborious/utils/repository/opc_repository.py @@ -2,13 +2,7 @@ from pathlib import Path from asyncua.sync import Client from asyncua.crypto.security_policies import SecurityPolicyBasic256 from asyncua.ua import DataValue, Variant, VariantType -from datetime import datetime -from time import sleep from logging import Logger -from statistics import mean, median -from typing import Callable - -from sientia_do.notifications.models import NotificationLevel data_type_map = { 'float': VariantType.Float, diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 8591f8a..605b69d 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -11,7 +11,7 @@ class FormatAndExportPrediction(): async def run(self, input_data: dict[str, Any]): path_flag = input_data['path_flag'] data = input_data['data'] - confidence = input_data['confidence'] + prediction_confidence = input_data['prediction_confidence'] if path_flag is None: # proceed with formatting and exporting @@ -21,7 +21,7 @@ class FormatAndExportPrediction(): 'data': data, 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], - 'prediction_confidence': confidence, + 'prediction_confidence': prediction_confidence, } ) @@ -32,7 +32,7 @@ class FormatAndExportPrediction(): { 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], - 'prediction_confidence': confidence, + 'prediction_confidence': prediction_confidence, 'comment': input_data['comment'] } ) diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 5b40362..60871fe 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -13,8 +13,14 @@ class PredictionProcess(): @workflow.run async def run(self, input_data: dict[str, Any]): data = input_data['data'] + schema = input_data['schema'] + table_name = input_data['table_name'] + model = input_data['model'] + filters = input_data['filters'] + model_name = input_data['model_name'] + model_retention = input_data['model_retention'] - path_flag, _confidence = await workflow.execute_activity_method( + path_flag, confidence = await workflow.execute_activity_method( Gates.input_gate, { 'filters': input_data['filters'], @@ -22,38 +28,102 @@ class PredictionProcess(): } ) - if path_flag == 'stop': - return - - if path_flag == 'continue': - # repeat last prediction - await workflow.execute_activity_method( - Postgres.repeat_last_prediction, - { - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'model': input_data['model'] - } - ) + if await self.path_flag_handler( + data, path_flag, confidence, schema, table_name, model + ): return response_data, last_timestamp = await workflow.execute_activity_method( MLFlow.transform_data, { 'data': data, - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] + 'model_name': model_name, + 'model_retention': model_retention } ) path_flag, confidence = await workflow.execute_activity_method( Gates.mlflow_gate, { - 'filters': input_data['filters'], + 'filters': filters, 'data': response_data, 'type': 'transform' } ) - if path_flag == 'stop': + if await self.path_flag_handler( + data, path_flag, confidence, schema, table_name, model + ): return + + response_data = await workflow.execute_activity_method( + MLFlow.request_predict, + { + 'data': response_data, + 'model_name': model_name, + 'model_retention': model_retention + } + ) + + path_flag, confidence = await workflow.execute_activity_method( + Gates.mlflow_gate, + { + 'filters': filters, + 'data': response_data, + 'type': 'predict' + } + ) + + if await self.path_flag_handler( + data, path_flag, confidence, schema, table_name, model + ): + return + + await workflow.execute_child_workflow( + 'format_and_export_prediction', + { + 'path_flag': path_flag, + 'data': response_data['content'], + 'prediction_confidence': confidence, + 'timestamp': response_data['timestamp'], + 'model_id': model, + 'model_name': model_name, + 'model_retention': model_retention + } + ) + + async def path_flag_handler(self, data: dict[str, Any], path_flag: str, + confidence: int, schema: str, table_name: str, + model: str): + if path_flag == 'stop': + return True + + elif path_flag == 'repeat': + # repeat last prediction + await workflow.execute_activity_method( + Postgres.repeat_last_prediction, + { + 'schema': schema, + 'table_name': table_name, + 'model': model + } + ) + return True + + elif path_flag == 'continue': + # call write workflow + workflow.execute_child_workflow( + 'format_and_export_prediction', + { + 'path_flag': path_flag, + 'data': data, + 'prediction_confidence': confidence, + 'timestamp': last_timestamp, + 'model_id': model, + 'model_name': model_name, + 'model_retention': model_retention + } + ) + return True + + return False diff --git a/requirements.txt b/requirements.txt index a03f80a..1ccf40c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,24 +1,5 @@ temporalio psycopg2-binary asyncua -mlflow==2.10.1 -scikit-learn==1.4.2 -scipy==1.13.0 -shap==0.46.0 -catboost==1.2.5 -hyperopt==0.2.7 -kaleido==0.2.1 -xgboost==2.0.2 -cloudpickle==3.0.0 -pathspec -kafka-python -sshtunnel git+ssh://git@github.com/Aignosi/sientia-dataops-library.git git+ssh://git@github.com/Aignosi/sientia-mlops-library.git -async-timeout -kubernetes -boto3 -scikit-optimize==0.10.2 -pandas==2.2.2 -PyYAML==6.0.1 -numpy==1.26.4 \ No newline at end of file diff --git a/tests/laborious/activities/test_base.py b/tests/laborious/activities/test_base.py new file mode 100644 index 0000000..9b704ae --- /dev/null +++ b/tests/laborious/activities/test_base.py @@ -0,0 +1,32 @@ +from unittest.mock import MagicMock +from laborious.activities.base import BaseActivity +from pytest import fixture +from sientia_do.notifications.models import Notification + + +@fixture +def base_activity(): + return BaseActivity( + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + +def test_prepare_activity(base_activity): + base_activity.notification_handler.base_notification = Notification( + project="project", + pipeline="pipeline", + trigger="-", + model_name="-", + model_id="-", + ) + + base_activity.prepare_activity( + schedule_name="test_schedule", + model_name="test_model", + model_id="test_model_id", + ) + + assert base_activity.notification_handler.base_notification.schedule_name == "test_schedule" + assert base_activity.notification_handler.base_notification.model_name == "test_model" + assert base_activity.notification_handler.base_notification.model_id == "test_model_id" diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 8edca83..ff84955 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -15,9 +15,9 @@ def gates(): @mark.asyncio -@patch('laborious.activities.gates.filter_functions') +@patch('laborious.activities.gates.input_filter_functions') async def test_input_gate_specific_variables_null_values_with_stop_policy_only( - filter_functions_mock, + input_filter_functions_mock, gates ): specific_variables_null_values_mock = MagicMock(return_value=True) @@ -26,9 +26,15 @@ async def test_input_gate_specific_variables_null_values_with_stop_policy_only( def functions_side_effect(x): if x == 'SPECIFIC_VARIABLES_NULL_VALUES': return specific_variables_null_values_mock + if x == 'path_confidence': + return { + 'stop': -1, + 'continue': 2, + 'repeat': -1 + } return empty_data_mock - filter_functions_mock.__getitem__.side_effect = functions_side_effect + input_filter_functions_mock.__getitem__.side_effect = functions_side_effect input_data = { 'filters': { @@ -40,7 +46,8 @@ async def test_input_gate_specific_variables_null_values_with_stop_policy_only( 'data': { 'variable': ['variable1', 'variable2'], 'value': [1, 2] - } + }, + 'path_priority': ['stop', 'continue', 'repeat'] } result = await gates.input_gate(input_data) @@ -55,9 +62,9 @@ async def test_input_gate_specific_variables_null_values_with_stop_policy_only( @mark.asyncio -@patch('laborious.activities.gates.filter_functions') +@patch('laborious.activities.gates.input_filter_functions') async def test_input_gate_specific_variables_null_values_with_continue_policy_only( - filter_functions_mock, + input_filter_functions_mock, gates ): specific_variables_null_values_mock = MagicMock(return_value=True) @@ -66,9 +73,15 @@ async def test_input_gate_specific_variables_null_values_with_continue_policy_on def functions_side_effect(x): if x == 'SPECIFIC_VARIABLES_NULL_VALUES': return specific_variables_null_values_mock + if x == 'path_confidence': + return { + 'stop': -1, + 'continue': 2, + 'repeat': -1 + } return empty_data_mock - filter_functions_mock.__getitem__.side_effect = functions_side_effect + input_filter_functions_mock.__getitem__.side_effect = functions_side_effect input_data = { 'filters': { @@ -80,7 +93,8 @@ async def test_input_gate_specific_variables_null_values_with_continue_policy_on 'data': { 'variable': ['variable1', 'variable2'], 'value': [1, 2] - } + }, + 'path_priority': ['stop', 'continue', 'repeat'] } result = await gates.input_gate(input_data) @@ -95,9 +109,9 @@ async def test_input_gate_specific_variables_null_values_with_continue_policy_on @mark.asyncio -@patch('laborious.activities.gates.filter_functions') +@patch('laborious.activities.gates.input_filter_functions') async def test_input_gate_specific_variables_null_values_no_filtered( - filter_functions_mock, + input_filter_functions_mock, gates ): specific_variables_null_values_mock = MagicMock(return_value=False) @@ -106,9 +120,15 @@ async def test_input_gate_specific_variables_null_values_no_filtered( def functions_side_effect(x): if x == 'SPECIFIC_VARIABLES_NULL_VALUES': return specific_variables_null_values_mock + if x == 'path_confidence': + return { + 'stop': -1, + 'continue': 2, + 'repeat': -1 + } return empty_data_mock - filter_functions_mock.__getitem__.side_effect = functions_side_effect + input_filter_functions_mock.__getitem__.side_effect = functions_side_effect input_data = { 'filters': { @@ -120,7 +140,8 @@ async def test_input_gate_specific_variables_null_values_no_filtered( 'data': { 'variable': ['variable1', 'variable2'], 'value': [1, 2] - } + }, + 'path_priority': ['stop', 'continue', 'repeat'] } result = await gates.input_gate(input_data) @@ -137,9 +158,9 @@ async def test_input_gate_specific_variables_null_values_no_filtered( @mark.asyncio -@patch('laborious.activities.gates.filter_functions') +@patch('laborious.activities.gates.input_filter_functions') async def test_input_gate_one_stop_policy( - filter_functions_mock, + input_filter_functions_mock, gates ): specific_variables_null_values_mock = MagicMock(return_value=True) @@ -148,9 +169,15 @@ async def test_input_gate_one_stop_policy( def functions_side_effect(x): if x == 'SPECIFIC_VARIABLES_NULL_VALUES': return specific_variables_null_values_mock + if x == 'path_confidence': + return { + 'stop': -1, + 'continue': 2, + 'repeat': -1 + } return empty_data_mock - filter_functions_mock.__getitem__.side_effect = functions_side_effect + input_filter_functions_mock.__getitem__.side_effect = functions_side_effect input_data = { 'filters': { @@ -165,7 +192,8 @@ async def test_input_gate_one_stop_policy( 'data': { 'variable': ['variable1', 'variable2'], 'value': [1, 2] - } + }, + 'path_priority': ['stop', 'continue', 'repeat'] } result = await gates.input_gate(input_data) @@ -184,9 +212,9 @@ async def test_input_gate_one_stop_policy( @mark.asyncio -@patch('laborious.activities.gates.filter_functions') +@patch('laborious.activities.gates.input_filter_functions') async def test_input_gate_one_continue_policy( - filter_functions_mock, + input_filter_functions_mock, gates ): specific_variables_null_values_mock = MagicMock(return_value=False) @@ -195,9 +223,15 @@ async def test_input_gate_one_continue_policy( def functions_side_effect(x): if x == 'SPECIFIC_VARIABLES_NULL_VALUES': return specific_variables_null_values_mock + if x == 'path_confidence': + return { + 'stop': -1, + 'continue': 2, + 'repeat': -1 + } return empty_data_mock - filter_functions_mock.__getitem__.side_effect = functions_side_effect + input_filter_functions_mock.__getitem__.side_effect = functions_side_effect input_data = { 'filters': { @@ -212,7 +246,8 @@ async def test_input_gate_one_continue_policy( 'data': { 'variable': ['variable1', 'variable2'], 'value': [1, 2] - } + }, + 'path_priority': ['stop', 'continue', 'repeat'] } result = await gates.input_gate(input_data) @@ -231,9 +266,9 @@ async def test_input_gate_one_continue_policy( @mark.asyncio -@patch('laborious.activities.gates.filter_functions') +@patch('laborious.activities.gates.input_filter_functions') async def test_input_gate_no_filtered( - filter_functions_mock, + input_filter_functions_mock, gates ): specific_variables_null_values_mock = MagicMock(return_value=False) @@ -242,9 +277,15 @@ async def test_input_gate_no_filtered( def functions_side_effect(x): if x == 'SPECIFIC_VARIABLES_NULL_VALUES': return specific_variables_null_values_mock + if x == 'path_confidence': + return { + 'stop': -1, + 'continue': 2, + 'repeat': -1 + } return empty_data_mock - filter_functions_mock.__getitem__.side_effect = functions_side_effect + input_filter_functions_mock.__getitem__.side_effect = functions_side_effect input_data = { 'filters': { @@ -259,7 +300,8 @@ async def test_input_gate_no_filtered( 'data': { 'variable': ['variable1', 'variable2'], 'value': [1, 2] - } + }, + 'path_priority': ['stop', 'continue', 'repeat'] } result = await gates.input_gate(input_data) @@ -276,12 +318,12 @@ async def test_input_gate_no_filtered( @mark.asyncio -@patch('laborious.activities.gates.filter_functions') +@patch('laborious.activities.gates.input_filter_functions') async def test_input_gate_error( - filter_functions_mock, + input_filter_functions_mock, gates ): - filter_functions_mock.__getitem__.side_effect = KeyError('test') + input_filter_functions_mock.__getitem__.side_effect = KeyError('test') input_data = { 'filters': { @@ -293,7 +335,8 @@ async def test_input_gate_error( 'data': { 'variable': ['variable1', 'variable2'], 'value': [1, 2] - } + }, + 'path_priority': ['stop', 'continue', 'repeat'], } result = await gates.input_gate(input_data) @@ -306,3 +349,239 @@ async def test_input_gate_error( level=NotificationLevel.ERROR, attachment_content=ANY ) + + +transform_filter_path_confidence = { + 'stop': -1, + 'continue': 255, + 'repeat': -1 +} + + +@mark.asyncio +@patch('laborious.activities.gates.mlflow_response_filter_functions') +async def test_mlflow_response_gate_no_filtered( + mlflow_response_filter_functions_mock, + gates +): + api_error_filter_mock = MagicMock(return_value=False) + + def transform_filter_functions_side_effect(x: str): + if x == 'API_ERROR': + return api_error_filter_mock + if x == 'path_confidence': + return transform_filter_path_confidence + + mlflow_response_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect + + input_data = { + 'filters': { + 'API_ERROR': { + 'POLICY': 'stop', + } + }, + 'data': { + 'variable': ['variable1', 'variable2'], + 'value': [1, 2] + }, + 'path_priority': ['stop', 'continue', 'repeat'], + 'type': 'predict' + } + + result = await gates.mlflow_response_gate(input_data) + assert result == (None, 0) + + api_error_filter_mock.assert_called_once_with( + input_data['data'], + input_data['filters']['API_ERROR'] + ) + + gates.notification_handler.build_and_send_notification.assert_not_called() + + +@mark.asyncio +@patch('laborious.activities.gates.mlflow_response_filter_functions') +async def test_mlflow_response_gate_filtered( + mlflow_response_filter_functions_mock, + gates +): + api_error_filter_mock = MagicMock(return_value=True) + + def transform_filter_functions_side_effect(x: str): + if x == 'API_ERROR': + return api_error_filter_mock + if x == 'path_confidence': + return transform_filter_path_confidence + + mlflow_response_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect + + input_data = { + 'filters': { + 'API_ERROR': { + 'POLICY': 'continue', + } + }, + 'data': { + 'success': False, + 'content': { + 'message': 'Error', + 'traceback': 'Error' + } + }, + 'path_priority': ['stop', 'continue', 'repeat'], + 'type': 'predict' + } + + result = await gates.mlflow_response_gate(input_data) + assert result == ('continue', 255) + + api_error_filter_mock.assert_called_once_with( + input_data['data'], + input_data['filters']['API_ERROR'] + ) + + gates.notification_handler.build_and_send_notification.assert_called_once_with( + notification_id='PREDICT_GATE_RESPONSE_FILTER__API_ERROR', + message=input_data['data']['content']['message'], + block='mlflow_gate', + level=NotificationLevel.WARNING, + attachment_content=input_data['data']['content']['traceback'] + ) + + +@mark.asyncio +@patch('laborious.activities.gates.mlflow_content_filter_functions') +async def test_mlflow_content_gate_no_filtered( + mlflow_content_filter_functions_mock, + gates +): + nan_values_filter_mock = MagicMock(return_value=False) + + def transform_filter_functions_side_effect(x: str): + if x == 'NAN_VALUES': + return nan_values_filter_mock + if x == 'path_confidence': + return transform_filter_path_confidence + + mlflow_content_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect + + input_data = { + 'filters': { + 'NAN_VALUES': { + 'POLICY': 'repeat', + } + }, + 'data': { + 'variable': ['variable1', 'variable2'], + 'value': [1, 2] + }, + 'path_priority': ['stop', 'continue', 'repeat'], + 'type': 'predict' + } + + result = await gates.mlflow_content_gate(input_data) + assert result == (None, 0) + + nan_values_filter_mock_args = nan_values_filter_mock.call_args + assert nan_values_filter_mock_args[0][0].equals(DataFrame( + {'variable': ['variable1', 'variable2'], 'value': [1, 2]})) + assert nan_values_filter_mock_args[0][1] == input_data['filters']['NAN_VALUES'] + + gates.notification_handler.build_and_send_notification.assert_not_called() + + +@mark.asyncio +@patch('laborious.activities.gates.mlflow_content_filter_functions') +async def test_mlflow_content_gate_filtered( + mlflow_content_filter_functions_mock, + gates +): + nan_values_filter_mock = MagicMock(return_value=True) + + def transform_filter_functions_side_effect(x: str): + if x == 'NAN_VALUES': + return nan_values_filter_mock + if x == 'path_confidence': + return transform_filter_path_confidence + + mlflow_content_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect + + input_data = { + 'filters': { + 'NAN_VALUES': { + 'POLICY': 'repeat', + } + }, + 'data': { + 'variable': ['variable1', 'variable2'], + 'value': [1, 2] + }, + 'path_priority': ['stop', 'continue', 'repeat'], + 'type': 'predict' + } + + result = await gates.mlflow_content_gate(input_data) + assert result == ('repeat', -1) + + nan_values_filter_mock_args = nan_values_filter_mock.call_args + assert nan_values_filter_mock_args[0][0].equals(DataFrame( + {'variable': ['variable1', 'variable2'], 'value': [1, 2]})) + assert nan_values_filter_mock_args[0][1] == input_data['filters']['NAN_VALUES'] + + gates.notification_handler.build_and_send_notification.assert_called_once_with( + notification_id='PREDICT_GATE_CONTENT_FILTER__NAN_VALUES', + message="Data not passed the content filter NAN_VALUES:{'POLICY': 'repeat'}", + block='mlflow_gate', + level=NotificationLevel.WARNING, + attachment_content=DataFrame(input_data['data']).to_string() + ) + + +@mark.asyncio +async def test_format_prediction( + gates +): + input_data = { + 'data': { + 'variable': ['variable1', 'variable2'], + 'value': [1, 2] + }, + 'timestamp': '2021-01-01', + 'model_id': 'model_id', + 'prediction_confidence': 0.95 + } + + expected_output = DataFrame(input_data['data']) + expected_output['timestamp'] = input_data['timestamp'] + expected_output['model_id'] = input_data['model_id'] + expected_output['prediction_confidence'] = input_data['prediction_confidence'] + expected_output['prediction_status'] = 'Good' + expected_output['comment'] = '' + + result = await gates.format_prediction(input_data) + assert result == expected_output.to_dict() + + +@mark.asyncio +async def test_format_default_prediction( + gates +): + input_data = { + 'timestamp': '2021-01-01', + 'model_id': 'model_id', + 'prediction_confidence': 0.95, + 'comment': 'Comment' + } + + expected_output = DataFrame({ + 'prediction': [0], + 'response_time': [0], + 'timestamp': [input_data['timestamp']], + 'model_id': [input_data['model_id']], + 'prediction_confidence': [input_data['prediction_confidence']], + 'prediction_status': ['Bad'], + 'comment': [input_data['comment']] + }) + + result = await gates.format_default_prediction(input_data) + assert result == expected_output.to_dict() diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py new file mode 100644 index 0000000..e470564 --- /dev/null +++ b/tests/laborious/activities/test_mlflow.py @@ -0,0 +1,121 @@ +from unittest.mock import MagicMock, patch + +import numpy as np +from pytest import fixture, mark +from laborious.activities.mlflow import MLFlow + + +@patch("laborious.activities.mlflow.MLFlowRepository") +def test___init__(mock_mlflow_repository): + mlflow = MLFlow( + mlflow_host="http://localhost", + mlflow_port=5000, + mlflow_username="admin", + mlflow_password="admin", + logger=MagicMock(), + notification_handler=MagicMock() + ) + + assert mlflow.mlflow_host == "http://localhost" + assert mlflow.mlflow_port == 5000 + assert mlflow.mlflow_username == "admin" + assert mlflow.mlflow_password == "admin" + + mock_mlflow_repository.assert_called_once_with( + "http://localhost:5000", "admin", "admin" + ) + + +@fixture +@patch("laborious.activities.mlflow.MLFlowRepository") +def mlflow(mock_mlflow_repository): + return MLFlow( + mlflow_host="http://localhost:5000", + mlflow_port=5000, + mlflow_username="admin", + mlflow_password="admin", + logger=MagicMock(), + notification_handler=MagicMock() + ) + + +@mark.asyncio +@patch("laborious.activities.mlflow.DataFrame") +@patch("laborious.activities.mlflow.max") +async def test_request_transform(mock_max, mock_dataframe, mlflow): + mock_max.return_value = '2024-01-02' + # Mock input data + input_data = { + 'data': [ + {'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0}, + {'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0}, + {'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0}, + {'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0} + ], + 'model_name': 'test_model', + 'model_retention': 30 + } + + # Mock the transform response + expected_response = {'prediction': [0.5, 0.6]} + mlflow.model_monitoring_repository.transform.return_value = expected_response + + # Call the method + response_data, timestamp = await mlflow.request_transform(input_data) + + # Verify the data was correctly transformed + mock_dataframe.assert_called_once_with(input_data['data']) + mock_dataframe.return_value.pivot.assert_called_once_with( + index='timestamp', columns='variable', values='value' + ) + mock_dataframe = mock_dataframe.return_value.pivot.return_value + mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True) + mock_dataframe.reset_index.assert_called_once() + mock_dataframe.columns.name = None + + # Verify the response + assert response_data == expected_response + assert timestamp == '2024-01-02' + + # Verify the repository was called with correct arguments + mlflow.model_monitoring_repository.transform.assert_called_once_with( + 'test_model', mock_dataframe, 30 + ) + + +@mark.asyncio +@patch("laborious.activities.mlflow.DataFrame") +@patch("laborious.activities.mlflow.max") +async def test_request_predict(mock_max, mock_dataframe, mlflow): + mock_max.return_value = '2024-01-02' + # Mock input data + input_data = { + 'data': [ + {'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0}, + {'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0}, + {'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0}, + {'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0} + ], + 'model_name': 'test_model', + 'model_retention': 30 + } + + # Mock the predict response + expected_response = {'prediction': [0.5, 0.6]} + mlflow.model_monitoring_repository.predict.return_value = expected_response + + # Call the method + response_data = await mlflow.request_predict(input_data) + + mock_dataframe.assert_called_once_with(input_data['data']) + mock_dataframe.return_value.replace.assert_called_once_with( + np.nan, None, inplace=True + ) + + # Verify the response + assert response_data == expected_response + + # Verify the repository was called with correct arguments + mlflow.model_monitoring_repository.predict.assert_called_once_with( + 'test_model', mock_dataframe.return_value, 30 + ) diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py new file mode 100644 index 0000000..2cde0af --- /dev/null +++ b/tests/laborious/activities/test_opc.py @@ -0,0 +1,178 @@ +from unittest.mock import patch, MagicMock + +from pytest import fixture, mark +from laborious.activities.opc import OPC +from sientia_do.notifications.models import NotificationLevel +from unittest.mock import ANY + + +@patch("laborious.activities.opc.OpcRepository") +def test___init__(mock_opc_repository): + opc = OPC( + name="test", + url="http://localhost:8080", + server_uri="opc.tcp://localhost:4840", + cert_path="", + private_key_path="", + server_cert_path="", + logger=MagicMock(), + notification_handler=MagicMock() + ) + + assert opc.name == "test" + assert opc.url == "http://localhost:8080" + assert opc.server_uri == "opc.tcp://localhost:4840" + assert opc.cert_path == "" + assert opc.private_key_path == "" + assert opc.server_cert_path == "" + assert opc.opc_repository == mock_opc_repository.return_value + + mock_opc_repository.assert_called_once_with( + name="test", + url="http://localhost:8080", + server_uri="opc.tcp://localhost:4840", + cert_path="", + private_key_path="", + server_cert_path="", + logger=opc.logger, + ) + + opc.opc_repository.connect.assert_called_once() + + +@fixture +@patch("laborious.activities.opc.OpcRepository") +def opc(mock_opc_repository): + return OPC( + name="test", + url="http://localhost:8080", + server_uri="opc.tcp://localhost:4840", + cert_path="", + private_key_path="", + server_cert_path="", + logger=MagicMock(), + notification_handler=MagicMock() + ) + + +@mark.asyncio +async def test_write_opc_data_success(opc): + # Arrange + input_data = { + 'data': { + 'prediction': [0.75], + 'prediction_confidence': [0.95] + }, + 'opc_servers': ['server1'], + 'opc_output_config': { + 'prediction_tags': { + 'tag1': {'data_type': 'float'} + }, + 'confidence_tags': { + 'tag2': {'data_type': 'float'} + } + } + } + + # Act + await opc.write_opc_data(input_data) + + # Assert + opc.opc_repository.write_data.assert_any_call('tag1', 0.75, 'float') + opc.opc_repository.write_data.assert_any_call('tag2', 0.95, 'float') + assert opc.opc_repository.write_data.call_count == 2 + + +@mark.asyncio +async def test_write_opc_data_prediction_error(opc): + # Arrange + input_data = { + 'data': { + 'prediction': [0.75], + 'prediction_confidence': [0.95] + }, + 'opc_servers': ['server1'], + 'opc_output_config': { + 'prediction_tags': { + 'tag1': {'data_type': 'float'} + } + } + } + + opc.opc_repository.write_data.side_effect = Exception("Test error") + + # Act + await opc.write_opc_data(input_data) + + # Assert + opc.notification_handler.build_and_send_notification.assert_called_with( + notification_id="WRITE_OPC_PREDICTION_ERROR", + message="Error writing data to OPC server: Test error", + block="write_opc_data", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + opc.logger.error.assert_called_once() + + +@mark.asyncio +async def test_write_opc_data_confidence_error(opc): + # Arrange + input_data = { + 'data': { + 'prediction': [0.75], + 'prediction_confidence': [0.95] + }, + 'opc_servers': ['server1'], + 'opc_output_config': { + 'prediction_tags': { + 'tag1': {'data_type': 'float'} + }, + 'confidence_tags': { + 'tag2': {'data_type': 'float'} + } + } + } + + # Make first call succeed but second fail + def side_effect(*args, **kwargs): + if args[0] == 'tag2': + raise ValueError("Test error") + return None + + opc.opc_repository.write_data.side_effect = side_effect + + # Act + await opc.write_opc_data(input_data) + + # Assert + opc.notification_handler.build_and_send_notification.assert_called_with( + notification_id="WRITE_OPC_CONFIDENCE_ERROR", + message="Error writing data to OPC server: Test error", + block="write_opc_data", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + opc.logger.error.assert_called_once() + + +@mark.asyncio +async def test_write_opc_data_empty_config(opc): + # Arrange + input_data = { + 'data': { + 'prediction': [0.75], + 'prediction_confidence': [0.95] + }, + 'opc_servers': ['server1'], + 'opc_output_config': { + 'prediction_tags': {}, + 'confidence_tags': {} + } + } + + # Act + await opc.write_opc_data(input_data) + + # Assert + opc.opc_repository.write_data.assert_not_called() diff --git a/tests/laborious/utils/filters/repository/test_model_repository.py b/tests/laborious/utils/filters/repository/test_model_repository.py new file mode 100644 index 0000000..abf4ab8 --- /dev/null +++ b/tests/laborious/utils/filters/repository/test_model_repository.py @@ -0,0 +1,278 @@ +from unittest.mock import ANY, MagicMock, patch +import numpy as np +from pandas import DataFrame +import pytest +from laborious.utils.repository.model_repository import MLFlowRepository + + +@pytest.fixture +def mlflow_repository(): + with patch('laborious.utils.repository.model_repository.ModelServing', autospec=True) as MockModelServing: + mock_instance = MockModelServing.return_value + mock_instance.get_transformed_data = MagicMock() + + repo = MLFlowRepository( + host='http://localhost:5000', + username='admin', + password='admin' + ) + return repo + + +def test_get_current_data_df(mlflow_repository): + current_data = { + 'prediction': [1, 3], + 'target': [1, 1], + } + mlflow_repository.model_serving.get_transformed_data.return_value = { + 'var1': [1, 2], + 'var2': [2, np.nan], + } + expected = DataFrame({ + 'var1': [1], + 'var2': [2], + 'prediction': [1], + 'target': [1], + }) + output = mlflow_repository.get_current_data_df(current_data, + 'model', 'target') + + mlflow_repository.model_serving.get_transformed_data.assert_called_once_with( + 'model', current_data, by='model') + + diff = output.compare(expected) + assert diff.empty + + +def test_get_artifact(mlflow_repository): + mlflow_repository.get_artifact( + 'destination', 'search_by', 'run_id', 'model', 'artifact' + ) + mlflow_repository.model_serving.get_artifact.assert_called_once_with( + destination='destination', + search_by='search_by', + run_id='run_id', + model_name='model', + artifact_name='artifact' + ) + + +def test_calculate_model_metrics(mlflow_repository): + mlflow_repository.model_serving.get_model_metrics.return_value = 'data' + real_data = 'real_data' + predictions = 'predictions' + flag = 'flag' + output = mlflow_repository.calculate_model_metrics( + real_data, predictions, flag + ) + mlflow_repository.model_serving.get_model_metrics.assert_called_once_with( + reference_data=None, + real_data=real_data, + predictions=predictions, + type_flag=flag + ) + assert output == 'data' + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_by_run_id(mlflow, mlflow_repository): + mlflow.get_run.return_value = MagicMock( + info=MagicMock( + experiment_id='0', + ) + ) + mlflow.get_experiment.return_value = MagicMock() + mlflow.get_experiment.return_value.name = 'test' + + output = mlflow_repository.get_experiment_by_run_id('0') + assert output == 'test' + mlflow.get_run.assert_called_once_with('0') + mlflow.get_experiment.assert_called_once_with('0') + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_next_run_name(mlflow, mlflow_repository): + mlflow.search_runs.return_value = [1, 2, 3] + output = mlflow_repository.get_next_run_name('run') + assert output == 'run-4' + mlflow.search_runs.assert_called_once_with( + experiment_names=['run'], + order_by=['start_time desc'], + ) + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_success(mlflow, mlflow_repository): + mlflow.get_experiment_by_name.return_value = MagicMock( + experiment_id='0') + + output = mlflow_repository.get_experiment('test') + + assert output == 0 + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_error(mlflow, mlflow_repository): + mlflow.get_experiment_by_name.return_value = None + + try: + mlflow_repository.get_experiment('test') + except ValueError as e: + assert str(e) == 'Experiment test not found' + else: + assert False + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_experiment_last_run(mlflow, mlflow_repository): + mlflow.search_runs.return_value = DataFrame({ + 'params.retrain': ['True', 'False', 'True', 'False'], + 'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'], + 'run_id': ['0', '1', '2', '3'], + }) + + output = mlflow_repository.get_experiment_last_run(0) + + mlflow.search_runs.assert_called_once_with( + experiment_ids=[0], + filter_string="", + output_format="pandas", + ) + + assert output == '2' + + +@patch('laborious.utils.repository.model_repository.mlflow') +def test_update_production_model_by_run_id(mlflow, mlflow_repository): + client_mock = MagicMock() + mlflow.tracking.MlflowClient.return_value = client_mock + + client_mock.get_registered_model.return_value = MagicMock( + latest_versions=[ + MagicMock(version='1'), + MagicMock(version='2'), + MagicMock(version='3'), + ] + ) + output = mlflow_repository.update_production_model_by_run_id('0', 'test') + + mlflow.register_model.assert_called_once_with( + "runs:/0/prediction_model", + 'test', + ) + + mlflow.tracking.MlflowClient.assert_called_once() + client_mock.get_registered_model.assert_called_once_with('test') + client_mock.transition_model_version_stage.assert_called_once_with( + name='test', + version='3', + stage='Production', + archive_existing_versions=True, + ) + + assert output == { + 'model_name': 'test', + 'version': '3', + 'mlflow_run_id': '0', + } + + +def test_update_production_model(mlflow_repository): + connector = mlflow_repository + + with patch.object(connector, 'get_experiment', + return_value='0') as get_experiment: + with patch.object(connector, 'get_experiment_last_run', + return_value='2') as get_experiment_last_run: + with patch.object(connector, 'update_production_model_by_run_id', + return_value={'model_name': 'test', 'version': '3', + 'mlflow_run_id': '0'}) as update_production_model_by_run_id: + + output = connector.update_production_model('0', 'test') + + get_experiment.assert_called_once_with('0') + get_experiment_last_run.assert_called_once_with('0') + update_production_model_by_run_id.assert_called_once_with( + '2', 'test') + + assert output == { + 'model_name': 'test', + 'version': '3', + 'mlflow_run_id': '0', + 'mlflow_experiment_id': '0', + } + + +def test_transform_success(mlflow_repository): + data = 'data' + model_name = 'model' + + output = mlflow_repository.transform(model_name, data, 1) + + mlflow_repository.model_serving.get_cached_transform.assert_called_once_with( + model_name, data, 1) + + assert output == { + 'success': True, + 'content': mlflow_repository.model_serving.get_cached_transform.return_value.to_dict.return_value + } + + +def test_transform_error(mlflow_repository): + data = 'data' + model_name = 'model' + + mlflow_repository.model_serving.get_cached_transform.side_effect = Exception( + 'error') + + output = mlflow_repository.transform(model_name, data, 1) + + mlflow_repository.model_serving.get_cached_transform.assert_called_once_with( + model_name, data, 1) + + assert output == { + 'success': False, + 'content': { + 'message': 'error', + 'traceback': ANY + } + } + + +def test_predict_success(mlflow_repository): + data = 'data' + model_name = 'model' + mlflow_repository.model_serving.get_cached_predict.return_value = np.array( + [2, 3] + ) + + output = mlflow_repository.predict(model_name, data, 1) + + mlflow_repository.model_serving.get_cached_predict.assert_called_once_with( + model_name, data, 1) + + assert output['success'] == True + assert output['content'] == {'prediction': { + 0: 2, 1: 3}, 'response_time': ANY} + + +def test_predict_error(mlflow_repository): + data = 'data' + model_name = 'model' + + mlflow_repository.model_serving.get_cached_predict = MagicMock( + side_effect=Exception('error') + ) + + output = mlflow_repository.predict(model_name, data, 1) + + mlflow_repository.model_serving.get_cached_predict.assert_called_once_with( + model_name, data, 1) + + assert output == { + 'success': False, + 'content': { + 'message': 'error', + 'traceback': ANY + } + } diff --git a/tests/laborious/utils/filters/repository/test_opc_repository.py b/tests/laborious/utils/filters/repository/test_opc_repository.py new file mode 100644 index 0000000..070f001 --- /dev/null +++ b/tests/laborious/utils/filters/repository/test_opc_repository.py @@ -0,0 +1,106 @@ +from unittest.mock import Mock, patch, MagicMock +from pathlib import Path +from asyncua.sync import Client +from asyncua.crypto.security_policies import SecurityPolicyBasic256 +from asyncua.ua import DataValue, Variant, VariantType +from pytest import fixture +from laborious.utils.repository.opc_repository import OpcRepository + + +@fixture +def mock_logger(): + return Mock() + + +@fixture +def opc_repository(mock_logger): + return OpcRepository( + name="test_repo", + url="opc.tcp://localhost:4840", + logger=mock_logger, + server_uri="urn:test:server", + cert_path="/path/to/cert.pem", + private_key_path="/path/to/key.pem", + server_cert_path="/path/to/server_cert.pem" + ) + + +@fixture +def mock_client(): + with patch('laborious.utils.repository.opc_repository.Client') as mock: + client_instance = MagicMock() + mock.return_value = client_instance + yield client_instance + + +def test_init(opc_repository): + assert opc_repository.name == "test_repo" + assert opc_repository.url == "opc.tcp://localhost:4840" + assert opc_repository.server_uri == "urn:test:server" + assert opc_repository.cert_path == "/path/to/cert.pem" + assert opc_repository.private_key_path == "/path/to/key.pem" + assert opc_repository.server_cert_path == "/path/to/server_cert.pem" + assert opc_repository.non_receive_count == 0 + assert opc_repository.client is None + + +def test_set_security(opc_repository, mock_client): + opc_repository.client = mock_client + opc_repository.set_security() + + mock_client.application_uri = "urn:test:server" + mock_client.set_security.assert_called_once_with( + SecurityPolicyBasic256, + certificate="/path/to/cert.pem", + private_key="/path/to/key.pem", + server_certificate="/path/to/server_cert.pem" + ) + assert mock_client.secure_channel_timeout == 10000000 + assert mock_client.session_timeout == 10000000 + + +def test_set_security_missing_certificates(opc_repository): + opc_repository.cert_path = None + opc_repository.private_key_path = None + + try: + opc_repository.set_security() + except ValueError as e: + assert str( + e) == "Certificate and private key paths must be provided for secure connection." + + +def test_connect_with_security(opc_repository, mock_client): + opc_repository.connect() + + mock_client.connect.assert_called_once() + assert opc_repository.client == mock_client + + +def test_connect_without_security(opc_repository, mock_client): + opc_repository.cert_path = None + opc_repository.connect() + + mock_client.connect.assert_called_once() + assert opc_repository.client == mock_client + + +def test_disconnect(opc_repository, mock_client): + opc_repository.client = mock_client + opc_repository.disconnect() + + mock_client.disconnect.assert_called_once() + assert opc_repository.client is None + + +def test_write_data(opc_repository, mock_client, mock_logger): + opc_repository.client = mock_client + mock_node = MagicMock() + mock_client.get_node.return_value = mock_node + + opc_repository.write_data("ns=2;s=TestNode", 42.0, "float", mock_logger) + + mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") + mock_node.write_value.assert_called_once() + mock_logger.info.assert_called_once_with( + "Writing 42.0 - to " + str(mock_node)) diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py index 5e9b7a0..acd073a 100644 --- a/tests/laborious/utils/filters/test_conditional_filters.py +++ b/tests/laborious/utils/filters/test_conditional_filters.py @@ -7,20 +7,21 @@ def test_filter_specific_variables_null_values(): assert filter_specific_variables_null_values( DataFrame( {'variable': ['variable1', 'variable2'], 'value': [1, 2]}), - variables=['variable2']) == True + config={'VARIABLES': ['variable2']}) == True def test_filter_specific_variables_null_values_with_null_values(): assert filter_specific_variables_null_values( DataFrame( {'variable': ['variable1', 'variable2'], 'value': [1, None]}), - variables=['variable2']) == False + config={'VARIABLES': ['variable2']}) == False def test_filter_empty_data(): - assert filter_empty_data(DataFrame()) == True + assert filter_empty_data(DataFrame(), {}) == True def test_filter_empty_data_with_data(): assert filter_empty_data( - DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]})) == False + DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), + {}) == False diff --git a/tests/laborious/utils/filters/test_mlflow_filters.py b/tests/laborious/utils/filters/test_mlflow_filters.py new file mode 100644 index 0000000..f9c61e9 --- /dev/null +++ b/tests/laborious/utils/filters/test_mlflow_filters.py @@ -0,0 +1,22 @@ +from pandas import DataFrame +from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter + + +def test_api_error_filter_invalid_response(): + assert api_error_filter(None, {}) == True # NOSONAR + + +def test_api_error_filter_valid_response_fail(): + assert api_error_filter({'success': False}, {}) == True + + +def test_api_error_filter_valid_response_success(): + assert api_error_filter({'success': True}, {}) == False + + +def test_nan_values_filter_all_nan_values(): + assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True + + +def test_nan_values_filter_no_nan_values(): + assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False diff --git a/tests/laborious/workflows/subworkflows.py/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows.py/test_format_and_export_prediction.py new file mode 100644 index 0000000..31e5db1 --- /dev/null +++ b/tests/laborious/workflows/subworkflows.py/test_format_and_export_prediction.py @@ -0,0 +1,115 @@ +from unittest.mock import call, patch, AsyncMock +from pytest import mark, fixture + +from laborious.activities.activities import Activities +from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction + + +@fixture +def format_and_export_prediction(): + return FormatAndExportPrediction() + + +@mark.asyncio +@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): + + input_data = { + "path_flag": None, + "data": {"test": "data"}, + "timestamp": "2021-01-01", + "model_id": 1, + "prediction_confidence": 0, + "schema": "test_schema", + "table_name": "test_table", + "opc_servers": ["test_server"], + "opc_output_config": {"test": "config"} + } + + await format_and_export_prediction.run(input_data) + + workflow_mock.execute_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'] + } + )]) + 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': workflow_mock.execute_activity_method.return_value + } + )]) + workflow_mock.execute_activity_method.assert_has_calls([ + call( + Activities.write_opc_data, + { + 'opc_servers': input_data['opc_servers'], + 'opc_output_config': input_data['opc_output_config'], + 'data': workflow_mock.execute_activity_method.return_value + } + ) + ]) + + assert workflow_mock.execute_activity_method.call_count == 3 + + +@mark.asyncio +@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +async def test_run_default_path_flag(workflow_mock, format_and_export_prediction): + + input_data = { + "path_flag": "default", + "data": {"test": "data"}, + "timestamp": "2021-01-01", + "model_id": 1, + "prediction_confidence": 0, + "schema": "test_schema", + "table_name": "test_table", + "opc_servers": ["test_server"], + "opc_output_config": {"test": "config"}, + "comment": "test_comment" + } + + await format_and_export_prediction.run(input_data) + + workflow_mock.execute_activity_method.assert_has_calls([ + call( + Activities.format_default_prediction, + { + 'timestamp': input_data['timestamp'], + 'model_id': input_data['model_id'], + 'prediction_confidence': input_data['prediction_confidence'], + 'comment': input_data['comment'] + } + ) + ]) + 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': workflow_mock.execute_activity_method.return_value + } + ) + ]) + workflow_mock.execute_activity_method.assert_has_calls([ + call( + Activities.write_opc_data, + { + 'opc_servers': input_data['opc_servers'], + 'opc_output_config': input_data['opc_output_config'], + 'data': workflow_mock.execute_activity_method.return_value + } + ) + ]) + + assert workflow_mock.execute_activity_method.call_count == 3