From 4f37b2be022cdadcdbfa0829a242a9993d2d7ca2 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 30 May 2025 16:24:52 -0300 Subject: [PATCH] SIENTIAPDE-1030 Add unit tests for orchestrator activities and workflows - Implement tests for Activities class, covering initialization and prepare_activity method. - Create tests for Couchbase class, including successful and failed query loading. - Add tests for SlotManager class, verifying OPC slot loading and active ingestor retrieval. - Develop tests for TemporalManager class, focusing on schedule loading functionality. - Introduce tests for Orchestrator class, ensuring proper execution of workflow activities. - Establish a new test suite for orchestrator activities and workflows in the tests directory. --- .gitignore | 2 + docker-compose.yml | 92 +- laborious/activities/activities.py | 53 -- laborious/activities/base.py | 26 - laborious/activities/gates.py | 297 ------ laborious/activities/mlflow.py | 91 -- laborious/activities/opc.py | 105 --- laborious/activities/postgres.py | 181 ---- .../utils/filters/conditional_filters.py | 16 - laborious/utils/filters/mlflow_filters.py | 22 - laborious/utils/logger.py | 22 - laborious/utils/policies.py | 9 - .../utils/repository/model_repository.py | 297 ------ laborious/utils/repository/opc_repository.py | 207 ----- laborious/workflows/predictions_batch.py | 89 -- .../format_and_export_prediction.py | 95 -- .../sub_workflows/prediction_process.py | 233 ----- {laborious => orchestrator}/__init__.py | 0 .../activities/__init__.py | 0 orchestrator/activities/activities.py | 47 + orchestrator/activities/couchbase.py | 94 ++ orchestrator/activities/formatters.py | 130 +++ .../activities/notification.py | 0 orchestrator/activities/slot_manager.py | 73 ++ orchestrator/activities/temporal_manager.py | 66 ++ .../utils}/__init__.py | 0 .../utils/connectors_config.py | 0 .../worker/__init__.py | 0 {laborious => orchestrator}/worker/worker.py | 0 .../workflows/__init__.py | 0 orchestrator/workflows/orchestrator.py | 69 ++ requirements.txt | 3 +- samples.json | 35 + simulator/Dockerfile | 30 - test.ipynb | 876 ++++++++++++++++++ tests/laborious/activities/__init__.py | 0 tests/laborious/activities/test_activities.py | 193 ---- tests/laborious/activities/test_base.py | 35 - tests/laborious/activities/test_gates.py | 369 -------- tests/laborious/activities/test_mlflow.py | 120 --- tests/laborious/activities/test_opc.py | 196 ---- tests/laborious/activities/test_postgres.py | 159 ---- tests/laborious/utils/__init__.py | 0 tests/laborious/utils/filters/__init__.py | 0 .../utils/filters/test_conditional_filters.py | 30 - .../utils/filters/test_mlflow_filters.py | 22 - .../utils/repository/test_model_repository.py | 278 ------ .../utils/repository/test_opc_repository.py | 259 ------ .../laborious/utils/test_connectors_config.py | 133 --- tests/laborious/utils/test_logger.py | 37 - .../test_format_and_export_prediction.py | 127 --- .../subworkflows/test_prediction_process.py | 515 ---------- .../workflows/test_predictions_batch.py | 81 -- tests/{laborious => orchestrator}/__init__.py | 0 .../activities/test_activities.py | 116 +++ .../orchestrator/activities/test_couchbase.py | 56 ++ .../activities/test_slot_manager.py | 57 ++ .../activities/test_temporal_manager.py | 80 ++ .../workflows/test_orchestrator.py | 81 ++ 59 files changed, 1827 insertions(+), 4377 deletions(-) delete mode 100644 laborious/activities/activities.py delete mode 100644 laborious/activities/base.py delete mode 100644 laborious/activities/gates.py delete mode 100644 laborious/activities/mlflow.py delete mode 100644 laborious/activities/opc.py delete mode 100644 laborious/activities/postgres.py delete mode 100644 laborious/utils/filters/conditional_filters.py delete mode 100644 laborious/utils/filters/mlflow_filters.py delete mode 100644 laborious/utils/logger.py delete mode 100644 laborious/utils/policies.py delete mode 100644 laborious/utils/repository/model_repository.py delete mode 100644 laborious/utils/repository/opc_repository.py delete mode 100644 laborious/workflows/predictions_batch.py delete mode 100644 laborious/workflows/sub_workflows/format_and_export_prediction.py delete mode 100644 laborious/workflows/sub_workflows/prediction_process.py rename {laborious => orchestrator}/__init__.py (100%) rename {laborious => orchestrator}/activities/__init__.py (100%) create mode 100644 orchestrator/activities/activities.py create mode 100644 orchestrator/activities/couchbase.py create mode 100644 orchestrator/activities/formatters.py rename laborious/utils/__init__.py => orchestrator/activities/notification.py (100%) create mode 100644 orchestrator/activities/slot_manager.py create mode 100644 orchestrator/activities/temporal_manager.py rename {laborious/utils/filters => orchestrator/utils}/__init__.py (100%) rename {laborious => orchestrator}/utils/connectors_config.py (100%) rename {laborious => orchestrator}/worker/__init__.py (100%) rename {laborious => orchestrator}/worker/worker.py (100%) rename {laborious => orchestrator}/workflows/__init__.py (100%) create mode 100644 orchestrator/workflows/orchestrator.py create mode 100644 samples.json delete mode 100644 simulator/Dockerfile create mode 100644 test.ipynb delete mode 100644 tests/laborious/activities/__init__.py delete mode 100644 tests/laborious/activities/test_activities.py delete mode 100644 tests/laborious/activities/test_base.py delete mode 100644 tests/laborious/activities/test_gates.py delete mode 100644 tests/laborious/activities/test_mlflow.py delete mode 100644 tests/laborious/activities/test_opc.py delete mode 100644 tests/laborious/activities/test_postgres.py delete mode 100644 tests/laborious/utils/__init__.py delete mode 100644 tests/laborious/utils/filters/__init__.py delete mode 100644 tests/laborious/utils/filters/test_conditional_filters.py delete mode 100644 tests/laborious/utils/filters/test_mlflow_filters.py delete mode 100644 tests/laborious/utils/repository/test_model_repository.py delete mode 100644 tests/laborious/utils/repository/test_opc_repository.py delete mode 100644 tests/laborious/utils/test_connectors_config.py delete mode 100644 tests/laborious/utils/test_logger.py delete mode 100644 tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py delete mode 100644 tests/laborious/workflows/subworkflows/test_prediction_process.py delete mode 100644 tests/laborious/workflows/test_predictions_batch.py rename tests/{laborious => orchestrator}/__init__.py (100%) create mode 100644 tests/orchestrator/activities/test_activities.py create mode 100644 tests/orchestrator/activities/test_couchbase.py create mode 100644 tests/orchestrator/activities/test_slot_manager.py create mode 100644 tests/orchestrator/activities/test_temporal_manager.py create mode 100644 tests/orchestrator/workflows/test_orchestrator.py diff --git a/.gitignore b/.gitignore index f9bab0b..fd3c51f 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,8 @@ docker-compose.override.yml scouter/.file_versions/ scouter/pipelines/**/triggers.yaml **/postgres_data/** +**/couchbase_data/** +**/redis_data/** # Ignorar arquivos e diretórios de cache do Python __pycache__/ *.pyc diff --git a/docker-compose.yml b/docker-compose.yml index 532cee8..eb96070 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,63 +15,53 @@ services: networks: - sientia-network - zookeeper: - image: confluentinc/cp-zookeeper:7.5.1 - container_name: zookeeper - environment: - ZOOKEEPER_CLIENT_PORT: 2181 - ZOOKEEPER_TICK_TIME: 2000 + couchbase: + image: couchbase/server:7.2.0 + container_name: couchbase ports: - - "2181:2181" + - "8091:8091" # Admin UI and REST API + - "8092:8092" # Query Service (N1QL) + - "8093:8093" # Index Service + - "8094:8094" # Search Service + - "11210:11210" # Data Service (KV) + - "18091:18091" # Analytics Service (if enabled) + environment: + CB_CLUSTER_USERNAME: sientia + CB_CLUSTER_PASSWORD: sientia + CB_CLUSTER_RAMSIZE: 256 + CB_CLUSTER_INDEX_RAMSIZE: 256 + volumes: + - ./couchbase_data:/opt/couchbase/var + networks: + - sientia-network + healthcheck: + test: ["CMD-SHELL", "curl -f http://localhost:8091/pools/default || exit 1"] + interval: 10s + timeout: 10s + retries: 5 + + redis: + image: redis:7-alpine # Using a lightweight Redis image + container_name: redis + ports: + - "6379:6379" + volumes: + - ./redis_data:/data # Persist Redis data networks: - sientia-network - kafka: - image: confluentinc/cp-kafka:7.5.1 - container_name: kafka + redis-commander: + image: rediscommander/redis-commander:latest + container_name: redis-commander + environment: + REDIS_HOSTS: local:redis:6379 # Connects to the 'redis' service within the Docker network + ports: + - "8081:8081" # Access the Redis Commander UI on this port depends_on: - - zookeeper - ports: - - "9092:9092" - - "29092:29092" - environment: - KAFKA_BROKER_ID: 1 - KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181 - KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://kafka:29092,PLAINTEXT_HOST://localhost:9092 - KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: PLAINTEXT:PLAINTEXT,PLAINTEXT_HOST:PLAINTEXT - KAFKA_INTER_BROKER_LISTENER_NAME: PLAINTEXT - KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1 + - redis # Ensures Redis starts before Redis Commander networks: - sientia-network - kafka-ui: - image: provectuslabs/kafka-ui:latest - container_name: kafka-ui - ports: - - "8080:8080" - environment: - KAFKA_CLUSTERS_0_NAME: local - KAFKA_CLUSTERS_0_BOOTSTRAPSERVERS: kafka:29092 - networks: - - sientia-network - - simulator: - build: - context: . - dockerfile: simulator/Dockerfile - args: - GIT_REPO: ${SIMULATOR_GIT_REPO} - GIT_BRANCH: ${SIMULATOR_GIT_BRANCH} - container_name: simulator - ports: - - "4840:4840" - depends_on: - - kafka - networks: - - sientia-network - env_file: - - .env - networks: sientia-network: @@ -79,4 +69,8 @@ networks: volumes: postgres_data: + driver: local + couchbase_data: + driver: local + redis_data: driver: local \ No newline at end of file diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py deleted file mode 100644 index e42e847..0000000 --- a/laborious/activities/activities.py +++ /dev/null @@ -1,53 +0,0 @@ -from temporalio import activity, workflow - -with workflow.unsafe.imports_passed_through(): - from laborious.activities.postgres import Postgres - from laborious.activities.mlflow import MLFlow - from laborious.activities.gates import Gates - from laborious.activities.opc import OPC - from typing import Any - from logging import Logger - from sientia_do.notifications.handlers import NotificationHandler - - -class Activities(Postgres, MLFlow, Gates, OPC): - - def __init__(self, - postgres_config: dict[str, Any], - mlflow_config: dict[str, Any], - opc_config: dict[str, Any], - logger: Logger, notification_handler: NotificationHandler): - - # Initialize parent classes - Postgres.__init__(self, host=postgres_config['host'], - port=postgres_config['port'], - user=postgres_config['user'], - password=postgres_config['password'], - dbname=postgres_config['dbname'], - min_connections=postgres_config['min_connections'], - max_connections=postgres_config['max_connections'], - logger=logger, - notification_handler=notification_handler) - - MLFlow.__init__(self, mlflow_host=mlflow_config['host'], - mlflow_port=mlflow_config['port'], - mlflow_username=mlflow_config['username'], - mlflow_password=mlflow_config['password'], - logger=logger, - notification_handler=notification_handler) - - Gates.__init__(self, logger=logger, - notification_handler=notification_handler) - - OPC.__init__(self, - opc_servers=opc_config, - logger=logger, - notification_handler=notification_handler) - - @activity.defn(name="prepare_activity") - async def prepare_activity(self, input_data: dict[str, Any]): - await super().prepare_activity(input_data) - - def shutdown(self): - Postgres.close(self) - OPC.shutdown(self) diff --git a/laborious/activities/base.py b/laborious/activities/base.py deleted file mode 100644 index 3adb0e4..0000000 --- a/laborious/activities/base.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Any -from logging import Logger -from temporalio import activity -from sientia_do.notifications.handlers import NotificationHandler - - -class BaseActivity: - def __init__(self, logger: Logger, notification_handler: NotificationHandler): - self.logger = logger - self.notification_handler = notification_handler - - @activity.defn(name="prepare_activity") - async def prepare_activity(self, input_data: dict[str, Any]): - """ - Prepare the activity for the notification handler. - - Args: - workflow_name (str): The name of the workflow. - 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.pipeline_name = input_data['workflow_name'] - self.notification_handler.base_notification.schedule_name = input_data['schedule_name'] - self.notification_handler.base_notification.model_name = input_data['model_name'] - self.notification_handler.base_notification.model_id = input_data['model_id'] diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py deleted file mode 100644 index 1cf3fb9..0000000 --- a/laborious/activities/gates.py +++ /dev/null @@ -1,297 +0,0 @@ -from temporalio import activity, workflow - - -with workflow.unsafe.imports_passed_through(): - import traceback - from logging import Logger - from sientia_do.notifications.handlers import NotificationHandler - from sientia_do.notifications.models import NotificationLevel - from laborious.activities.base import BaseActivity - from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter - from typing import Any - from laborious.utils.filters.conditional_filters import ( - filter_empty_data, - filter_specific_variables_null_values - ) - from pandas import DataFrame - from datetime import datetime - -input_filter_functions = { - 'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values, - 'EMPTY_DATA': filter_empty_data, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 2, - 'REPEAT': -1 - } -} - -mlflow_response_filter_functions = { - 'API_ERROR': api_error_filter, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 10, - 'REPEAT': -1 - }, -} - -mlflow_content_filter_functions = { - 'NAN_VALUES': nan_values_filter, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 18, - 'REPEAT': -1 - } -} - - -class Gates(BaseActivity): - def __init__(self, logger: Logger, notification_handler: NotificationHandler): - BaseActivity.__init__(self, logger, notification_handler) - - @activity.defn(name="input_gate") - async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]: - """ - Filters the data based on the 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 filters to apply. - The key is the filter name and the value is the filter configuration. - data (dict[str, Any]): The data to filter. - path_priority (list[str]): The path priority. - Returns: - tuple[str | None, int, str]: (policy, confidence) based in priority - list and filter configuration and functions. - """ - - self.logger.debug("Performing input gate...") - - filters = input_data['filters'] - data = DataFrame(input_data['data']) - path_priority = input_data['path_priority'] - - filter_output = [] - - self.logger.debug(f"Input data:\n {data}") - self.logger.debug(f"Filters: {filters}") - - for fil, config in filters.items(): - if fil not in input_filter_functions: - self.logger.error(f"Filter {fil} not found") - continue - try: - if input_filter_functions[fil](data, config): - self.logger.debug( - f"Data not passed the input filter {fil}:{config}") - filter_output.append(config['POLICY']) - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id=f"INTPUT_GATE_ERROR__{fil}", - message=f"Error in filter {fil}:{config}: \n {e}", - block="input_gate", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - for path_flag in path_priority: - if path_flag in filter_output: - self.logger.debug(f"Input gate result: {path_flag}") - return path_flag, input_filter_functions['path_confidence'][path_flag], \ - "Input data with bad quality" - - self.logger.debug("Nothing was filtered by the input gate") - return None, 0, "" - - @activity.defn(name="mlflow_response_gate") - async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]: - """ - 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 | None, int, str]: (policy, confidence) based in priority list - and filter configuration and functions. - """ - - self.logger.debug("Performing mlflow response gate...") - - filters = input_data['filters'] - data = input_data['data'] - gate_type = input_data['type'] - path_priority = input_data['path_priority'] - - filter_output = [] - - self.logger.debug(f"Input data:\n {data}") - self.logger.debug(f"Filters: {filters}") - - comments = [] - for fil, config in filters.items(): - if fil not in mlflow_response_filter_functions: - continue - try: - if mlflow_response_filter_functions[fil](data, config): - filter_output.append(config['POLICY']) - comments.append(data['content']['message']) - self.notification_handler.build_and_send_notification( - notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}", - message=data['content']['message'], - block="mlflow_gate", - level=NotificationLevel.WARNING, - attachment_content=data['content']['traceback'] - ) - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}", - message=f"Error in filter {fil}:{config}: \n {e}", - block="mlflow_gate", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - for path_flag in path_priority: - if path_flag in filter_output: - self.logger.debug(f"Mlflow response gate result: {path_flag}") - return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \ - ", ".join(comments) - - self.logger.debug("Nothing was filtered by the mlflow response gate") - return None, 0, "" - - @activity.defn(name="mlflow_content_gate") - async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]: - """ - 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 | None, int, str]: (policy, confidence) based in priority - list and filter configuration and functions. - """ - - self.logger.debug("Performing mlflow content gate...") - - filters = input_data['filters'] - data = DataFrame(input_data['data']) - gate_type = input_data['type'] - path_priority = input_data['path_priority'] - - filter_output = [] - - self.logger.debug(f"Input data:\n {data}") - self.logger.debug(f"Filters: {filters}") - - for fil, config in filters.items(): - if fil not in mlflow_content_filter_functions: - continue - try: - 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}", - message=f"Data not passed the content filter {fil}:{config}", - block="mlflow_gate", - level=NotificationLevel.WARNING, - attachment_content=data.to_string() - ) - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}", - message=f"Error in filter {fil}:{config}: \n {e}", - block="mlflow_gate", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - for path_flag in path_priority: - if path_flag in filter_output: - self.logger.debug(f"Mlflow content gate result: {path_flag}") - return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \ - "Transformed data not passed the content filter" - - self.logger.debug("Nothing was filtered by the mlflow content gate") - return None, 0, "" - - @activity.defn(name="format_prediction") - 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. - """ - self.logger.debug("Formatting prediction...") - - data = DataFrame(input_data['data']) - data['timestamp'] = input_data['timestamp'] - data['model_id'] = input_data['model_id'] - data['prediction_confidence'] = input_data['prediction_confidence'] - data['prediction_status'] = 'Good' - data['comments'] = "" - data.sort_values(by='timestamp', inplace=True) - - return data.to_dict() - - @activity.defn(name="format_default_prediction") - 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. - """ - - self.logger.debug("Formatting default prediction...") - - return 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'], - 'comments': [input_data['comment']] - }).to_dict() - - @activity.defn(name="get_last_timestamp") - async def get_last_timestamp(self, input_data: dict[str, Any]) -> str: - """ - Gets the last timestamp of the data. - Args: - input_data (dict): The input data. Contains: - data (dict[str, Any]): The data to get the last timestamp from. - Returns: - str: The last timestamp of the data. - """ - data = DataFrame(input_data['data']) - if data.empty: - return datetime.now().strftime('%Y-%m-%d %H:%M:%S') - return max(data['timestamp'].values.tolist()) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py deleted file mode 100644 index ae141a2..0000000 --- a/laborious/activities/mlflow.py +++ /dev/null @@ -1,91 +0,0 @@ -import numpy as np -from pandas import DataFrame -from temporalio import activity, workflow - - -with workflow.unsafe.imports_passed_through(): - from laborious.activities.base import BaseActivity - from laborious.utils.repository.model_repository import MLFlowRepository - from typing import Any - from logging import Logger - from sientia_do.notifications.handlers import NotificationHandler - - -class MLFlow(BaseActivity): - def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str, - mlflow_password: str, logger: Logger, notification_handler: NotificationHandler): - BaseActivity.__init__(self, logger, notification_handler) - self.mlflow_host = mlflow_host - self.mlflow_port = mlflow_port - self.mlflow_username = mlflow_username - self.mlflow_password = mlflow_password - - 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]) -> dict[str, Any]: - """ - 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 in minutes. - Returns: - dict[str, Any]: The transformed data. - """ - self.logger.info('Transforming data...') - data = DataFrame(input_data['data']) - model_name = input_data['model_name'] - model_retention = input_data['model_retention'] - - self.logger.debug("Raw input data:") - self.logger.debug(data) - - data = data.pivot( - index='timestamp', columns='variable', - values='value') - data.fillna(np.nan, inplace=True) - data.reset_index(inplace=True) - data.columns.name = None - - self.logger.debug("Processed input data:") - self.logger.debug(data) - - response_data = self.model_monitoring_repository.transform( - model_name, data, model_retention) - - self.logger.debug("Response data:") - self.logger.debug(response_data) - - return response_data - - @activity.defn(name="request_predict") - async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]: - """ - 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: - dict[str, Any]: The predicted data. - """ - self.logger.info('Predicting data...') - data = DataFrame(input_data['data']) - model_name = input_data['model_name'] - model_retention = input_data['model_retention'] - - self.logger.debug(data) - - data.replace(np.nan, None, inplace=True) - - response_data = self.model_monitoring_repository.predict( - model_name, data, model_retention) - - self.logger.debug(response_data) - - return response_data diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py deleted file mode 100644 index 0f02c43..0000000 --- a/laborious/activities/opc.py +++ /dev/null @@ -1,105 +0,0 @@ -from temporalio import activity, workflow - - -with workflow.unsafe.imports_passed_through(): - from logging import Logger - from sientia_do.notifications.handlers import NotificationHandler - from sientia_do.notifications.models import NotificationLevel - from laborious.activities.base import BaseActivity - from laborious.utils.repository.opc_repository import OpcRepository - from typing import Any - import traceback - from pandas import DataFrame - - -class OPC(BaseActivity): - def __init__(self, opc_servers: dict[str, dict[str, Any]], - logger: Logger, notification_handler: NotificationHandler): - - self.logger = logger - self.notification_handler = notification_handler - self.opc_servers = opc_servers - - self.opc_repository = {} - for name, server in opc_servers.items(): - self.opc_repository[name] = OpcRepository( - name=name, - url=server['url'], - logger=self.logger, - server_uri=server['server_uri'], - cert_path=server['cert_path'], - private_key_path=server['private_key_path'], - server_cert_path=server['server_cert_path'], - notification_handler=self.notification_handler, - reconnection_interval=server['reconnection_interval'], - ) - self.opc_repository[name].connect() - - BaseActivity.__init__(self, logger, notification_handler) - - def write_data(self, server: str, tag: str, data: Any, - data_type: str, tag_type: str): - try: - self.opc_repository[server].write_data( - tag, data, data_type) - self.logger.debug(f"Wrote {tag_type} to {tag}") - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR", - message=f"Error writing data to OPC server: {e}", - block="write_opc_data", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - self.logger.error(trace) - - @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_output_config (dict[str, Any]): The OPC writing configuration. - The keys are the OPC server names and the values contain: - 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: - """ - self.logger.debug("Writing data to OPC servers...") - data = DataFrame(input_data['data']) - opc_output_config = input_data['opc_output_config'] - self.logger.debug(data) - - for server, config in opc_output_config.items(): - if self.opc_repository.get(server) is None: - self.logger.error(f"OPC server {server} not found") - continue - - if 'prediction_tags' in config: - for tag, tag_config in config['prediction_tags'].items(): - self.write_data( - server=server, - tag=tag, - data=data.head(1)['prediction'].values[0], - data_type=tag_config['data_type'], - tag_type='prediction' - ) - if 'confidence_tags' in config: - for tag, tag_config in config['confidence_tags'].items(): - self.write_data( - server=server, - tag=tag, - data=data.head(1)['prediction_confidence'].values[0], - data_type=tag_config['data_type'], - tag_type='confidence' - ) - - def shutdown(self): - for opc in self.opc_repository.values(): - opc.disconnect() diff --git a/laborious/activities/postgres.py b/laborious/activities/postgres.py deleted file mode 100644 index 9fc0ff7..0000000 --- a/laborious/activities/postgres.py +++ /dev/null @@ -1,181 +0,0 @@ -import traceback -from temporalio import workflow, activity - -from laborious.activities.base import BaseActivity -with workflow.unsafe.imports_passed_through(): - from sqlalchemy import create_engine - from sqlalchemy.orm import sessionmaker - from sqlalchemy.pool import QueuePool - from psycopg2.pool import ThreadedConnectionPool - from pandas import read_sql_query, DataFrame - from logging import Logger - from sientia_do.notifications.handlers import NotificationHandler - from sientia_do.notifications.models import NotificationLevel - from typing import Any - - -class Postgres(BaseActivity): - def __init__(self, host: str, port: int, - user: str, password: str, dbname: str, - min_connections: int, max_connections: int, - logger: Logger, notification_handler: NotificationHandler): - self.host = host - self.port = port - self.user = user - self.password = password - self.dbname = dbname - - # Create SQLAlchemy engine with connection pooling - self.engine = create_engine( - f'postgresql://{user}:{password}@{host}:{port}/{dbname}', - poolclass=QueuePool, - pool_size=min_connections, - max_overflow=max_connections - min_connections, - pool_pre_ping=True - ) - self.session_factory = sessionmaker(bind=self.engine) - - BaseActivity.__init__(self, logger, notification_handler) - - def close(self): - self.engine.dispose() - - def __del__(self): - self.close() - - @activity.defn(name="load_custom_query") - async def load_custom_query(self, query: str) -> dict[str, Any]: - """ - Loads data from a custom query. - - Args: - query (str): The query to load data from. - - Returns: - dict[str, dict]: The data from the query. - """ - self.logger.info(f"Fetching data from query: {query}") - - data = None - with self.session_factory() as session: - try: - data = read_sql_query(query, self.engine) - - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id="ERROR_LOADING_CUSTOM_QUERY", - message=f"Error fetching data from query: {e}", - block="load_custom_query", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - self.logger.error(trace) - - return {} - finally: - session.close() - - if data is None: - return {} - - # Converts any datetime datatype columns to string - for col in data.select_dtypes(include=['datetime64']).columns: - data[col] = data[col].dt.strftime('%Y-%m-%d %H:%M:%S') - - self.logger.info(f"Fetched {len(data)} rows") - self.logger.debug(f"Data: \n{data.to_string()}") - - return data.to_dict() - - @activity.defn(name="repeat_last_prediction") - async def repeat_last_prediction(self, query_items: dict[str, str]): - """ - Repeats the last prediction for a given model. - - Args: - query_items (dict[str, str]): The query items. Contains: - schema (str): The schema of the table. - table_name (str): The name of the table. - model (int): The model to repeat the prediction for. - - Returns: - None - """ - schema = query_items["schema"] - table_name = query_items["table_name"] - model = query_items["model"] - - repeat_query = f""" - INSERT INTO \"{schema}\".{table_name} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at) - SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW() - FROM \"{schema}\".{table_name} - WHERE model_id = {model} - ORDER BY timestamp DESC - LIMIT 1; - """ - self.logger.info(f"Repeating last prediction for model {model}") - self.logger.debug(f"Query: {repeat_query}") - - with self.session_factory() as session: - try: - session.execute(repeat_query) - session.commit() - - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id="ERROR_REPEATING_LAST_PREDICTION", - message=f"Error repeating last prediction: {e}", - block="repeat_last_prediction", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - self.logger.error(trace) - - finally: - session.close() - - @activity.defn(name="export_data_to_postgres") - async def export_data_to_postgres(self, input_data: dict[str, Any]): - """ - Exports data to a postgres table. - - Args: - 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. - """ - - self.logger.debug( - f"Exporting data to postgres: {input_data['data']}") - - schema = input_data["schema"] - table_name = input_data["table_name"] - data = DataFrame(input_data["data"]) - - with self.session_factory() as session: - try: - data.to_sql(table_name, self.engine, schema=schema, - if_exists="append", index=False) - session.commit() - - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES", - message=f"Error exporting data to postgres: {e}", - block="export_data_to_postgres", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - - self.logger.error(trace) - - else: - self.logger.debug("Data exported to postgres") - finally: - session.close() diff --git a/laborious/utils/filters/conditional_filters.py b/laborious/utils/filters/conditional_filters.py deleted file mode 100644 index 57cd3fd..0000000 --- a/laborious/utils/filters/conditional_filters.py +++ /dev/null @@ -1,16 +0,0 @@ -from pandas import DataFrame - - -def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool: - """ - Returns True if the specific columns have null values, False otherwise. - """ - return not data[ - data['variable'].isin(config['VARIABLES']) & data['value'].isna()].empty - - -def filter_empty_data(data: DataFrame, _config: dict) -> bool: - """ - Returns True if the data is empty, False otherwise. - """ - return data.empty diff --git a/laborious/utils/filters/mlflow_filters.py b/laborious/utils/filters/mlflow_filters.py deleted file mode 100644 index 9936018..0000000 --- a/laborious/utils/filters/mlflow_filters.py +++ /dev/null @@ -1,22 +0,0 @@ -import numpy as np -from pandas import DataFrame - - -def api_error_filter(response: dict, _config: dict): - if not response: - return True - - if not response['success']: - return True - - return False - - -def nan_values_filter(predictions: DataFrame, _config: dict): - data = predictions.replace({None: np.nan}).drop( - columns=['timestamp'], errors='ignore').infer_objects(copy=False) - - if data.isna().all().all(): - return True - - return False diff --git a/laborious/utils/logger.py b/laborious/utils/logger.py deleted file mode 100644 index 42a9cfd..0000000 --- a/laborious/utils/logger.py +++ /dev/null @@ -1,22 +0,0 @@ -from os import getenv -import logging -import sys - - -def get_logger(name: str): - log_level = getenv('LOG_LEVEL', 'INFO').upper() - - logger = logging.getLogger(name) - logger.setLevel(log_level) - stream_handler = logging.StreamHandler(sys.stdout) - stream_handler.setLevel(log_level) - - stream_handler.setFormatter( - logging.Formatter( - '%(asctime)s - %(name)s - %(levelname)s - %(message)s' - ) - ) - - logger.addHandler(stream_handler) - - return logger diff --git a/laborious/utils/policies.py b/laborious/utils/policies.py deleted file mode 100644 index 8c7449a..0000000 --- a/laborious/utils/policies.py +++ /dev/null @@ -1,9 +0,0 @@ -from datetime import timedelta -from temporalio.common import RetryPolicy - -retry_policy = RetryPolicy( - initial_interval=timedelta(seconds=1), - backoff_coefficient=2.0, - maximum_interval=timedelta(minutes=1), - maximum_attempts=1 -) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py deleted file mode 100644 index f2eaeb9..0000000 --- a/laborious/utils/repository/model_repository.py +++ /dev/null @@ -1,297 +0,0 @@ -""" -Model Monitoring Repository - -This module contains the ModelMonitoringRepository class, which is responsible for handling the communication with the Model Monitoring API. - -It includes the methods that are used to answer ModelMonitoringService requests using the Model Monitoring API functions. - -By Monitoring we mean the evaluation of the performance of models, the generation of reports. - -""" -from datetime import datetime -import traceback -import mlflow -import pandas as pd -from sientia.ModelServing import ModelServing - - -class MLFlowRepository(): - def __init__(self, host, username, password): - - self.model_serving = ModelServing(tracking_uri=host, - username=username, password=password) - - def get_current_data_df(self, current_data: pd.DataFrame, model_name: str, target: str): - """ - Get the current data as a DataFrame and update the prediction and target columns - - Parameters: - current_data (pd.DataFrame): the current data - model_name (str): the name of the model - target (str): the target column - - Returns: - DataFrame: the current data as a DataFrame - - - """ - predictions = current_data['prediction'] - - target = current_data[target] - current_data = self.model_serving.get_transformed_data( - model_name, current_data, by='model') - current_data['prediction'] = predictions - current_data['target'] = target - - return pd.DataFrame(current_data).dropna() - - def get_artifact(self, destination: str, search_by: str, run_id: str = None, - model_name: str = None, artifact_name: str = None) -> None: - """ - Get an artifact in MLflow by experiment or model and save it to a destination path using API. - If the artifact is searched by model, the latest production version will be used. - - Args: - destination: The destination path to save the artifact. - search_by: The way to search for the artifact ('experiment' or 'model'). - run_id: The run ID of the experiment (if search_by is "experiment"). - model_name: The name of the model (if search_by is "model"). - artifact_name: The path of the artifact to download. - - Returns: - artifact: The artifact(.csv) downloaded from MLflow. - """ - - self.model_serving.get_artifact(destination=destination, search_by=search_by, - run_id=run_id, model_name=model_name, artifact_name=artifact_name) - - def calculate_model_metrics(self, real_data, predictions, flag): - """ - Function to calculate the metrics of a model using API - - Parameters: - real_data (array): the real data - predictions (array): the predictions - - Returns: - dict: the metrics of the model including MSE and R2 - """ - return self.model_serving.get_model_metrics(reference_data=None, real_data=real_data, predictions=predictions, type_flag=flag) - - def get_experiment_by_run_id(self, run_id: str) -> dict: - # Get the run information using the run_id - run = mlflow.get_run(run_id) - - # Extract the experiment ID from the run - experiment_id = run.info.experiment_id - - # Get the experiment details using the experiment ID - experiment = mlflow.get_experiment(experiment_id) - experiment_name = experiment.name - return experiment_name - - def get_next_run_name(self, model_name: str) -> str: - """ - Function to get the next run number of a specific model - - Parameters: - model_name (str): the name of the model - - Returns: - str: the next run number - """ - - runs = mlflow.search_runs( - experiment_names=[model_name], order_by=["start_time desc"]) - next_run_number = len(runs) + 1 - return f"{model_name}-{next_run_number}" - - def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple: - """ - Retrain a model with new data. - - Parameters: - data (pandas.DataFrame): The new data to use for retraining. - model_name (str): The name of the model to retrain. - metrics_list (list): The metrics to be used to compare the models. - compare_metrics (bool): If True, the retrain will only be considered if the new model is better than the current one. - If False, the retrain will always be considered. - split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets. - If False, the data will be used as a unique block for retraining. - update_report (bool): If True, a report will be created with the data of the retrained model. - update_transformation (bool): If True, the model will be updated in the MLflow tracking server. - update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server. - shuffle_data (bool): If True, the data will be shuffled before splitting. - model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'. - - - Returns: - mlflow.sklearn.Model: The retrained prediction model. - mlflow.sklearn.Model: The retrained data model. - mse (float): The mean squared error of the retrained model. - r2 (float): The R-squared score of the retrained model. - """ - - # load predictor model - predictor_uri = f"models:/{model_name}/production" - # load transform model - latest_production_id = self.model_serving.get_model_run_id( - model_name, stage="Production" - ) - transform_uri = self.model_serving.get_model_uri( - latest_production_id, prediction=False - ) - # load - data_model = mlflow.sklearn.load_model(transform_uri) - prediction_model = mlflow.sklearn.load_model(predictor_uri) - data_model = data_model.fit(data) - treated_data = data_model.predict(data) - # align target column with treated_data - target_name = data_model.target_variable - y = data[target_name] - treated_data = pd.merge( - treated_data, y, left_index=True, right_index=True) - prediction_model = prediction_model.fit(treated_data) - # Example usage - experiment = self.get_experiment_by_run_id(latest_production_id) - pred_model_atributes = vars(prediction_model) # load class attributes - data_model_atributes = vars(data_model) # load class attributes - mlflow.set_experiment(experiment) - experiment_description = "Retrain model {model_name} with new data" - current_run_name = self.get_next_run_name(experiment) - with mlflow.start_run( - run_name=current_run_name, description=experiment_description - ) as _run: - # update transfomation model - # fixed parameters - for name_atribute, val_atribute in pred_model_atributes.items(): - if name_atribute != "model": - mlflow.log_param(name_atribute, val_atribute) - # update prediction model - for name_atribute, val_atribute in data_model_atributes.items(): - if name_atribute != "model": - mlflow.log_param(name_atribute, val_atribute) - # dynamic parameters, including model itself - mlflow.sklearn.log_model(data_model, "data_model") - file_path = f"laborious/data/raw_data_{model_name}.csv" - data.to_csv( - f"laborious/data/raw_data_{model_name}.csv", index=True) - # log the data raw - mlflow.log_artifact(file_path) - - # dynamic parameters, including model itself - mlflow.sklearn.log_model(prediction_model, "prediction_model") - mlflow.log_param("retrain", True) - - return "Model retrained successfully", experiment - - def get_experiment(self, experiment_name: str) -> int: - experiment = mlflow.get_experiment_by_name(experiment_name) - - if experiment is None: - raise ValueError(f'Experiment {experiment_name} not found') - - return int(experiment.experiment_id) - - def get_experiment_last_run(self, experiment_id: int) -> str: - runs = mlflow.search_runs( - experiment_ids=[experiment_id], - filter_string="", # Sem filtro no MLflow ainda - output_format="pandas" - ) - - # Filtrar apenas as runs onde params.retrain == True - filtered_runs = runs[runs["params.retrain"] == 'True'] - - # Converter a coluna 'end_time' para datetime - filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time']) - - # Ordenar o DataFrame de forma descendente pela coluna 'end_time' - filtered_runs = filtered_runs.sort_values( - by='end_time', ascending=False) - - # Pegar a última run_id do DataFrame filtrado e ordenado - latest_run_id = filtered_runs.iloc[0]['run_id'] - - return latest_run_id - - def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict: - # Registrar o modelo - # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro. - # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso. - mlflow.register_model( - f"runs:/{run_id}/prediction_model", model_name) - - # Colocar a versão do modelo em produção - # Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production' - client = mlflow.tracking.MlflowClient() - - # Obter a versão mais recente registrada do modelo - model_versions = client.get_registered_model( - model_name).latest_versions - max_version = max(model_versions, key=lambda x: int(x.version)).version - - # Mover a versão mais recente do modelo para o estágio de 'Production' - client.transition_model_version_stage( - name=model_name, - version=max_version, - stage="Production", - archive_existing_versions=True - ) - - return { - 'model_name': model_name, - 'version': max_version, - 'mlflow_run_id': run_id - } - - def update_production_model(self, experiment: str, model_name: str) -> dict: - - experiment_id = self.get_experiment(experiment) - run_id = self.get_experiment_last_run(experiment_id) - metadata = self.update_production_model_by_run_id(run_id, model_name) - - metadata['mlflow_experiment_id'] = experiment_id - - return metadata - - def transform(self, model_name: str, data: pd.DataFrame, model_retention: int): - try: - return { - 'success': True, - 'content': self.model_serving.get_cached_transform( - model_name, data, model_retention).to_dict() - } - - except Exception as e: - return { - 'success': False, - 'content': { - 'message': str(e), - 'traceback': traceback.format_exc() - } - } - - def predict(self, model_name: str, data: pd.DataFrame, model_retention: int): - try: - start_time = datetime.now() - data = self.model_serving.get_cached_predict( - model_name, data, model_retention)[-1:] - - end_time = datetime.now() - data = pd.DataFrame(data, columns=['prediction']) - data['response_time'] = (end_time - start_time).total_seconds() - - return { - 'success': True, - 'content': data.to_dict() - } - - except Exception as e: - return { - 'success': False, - 'content': { - 'message': str(e), - 'traceback': traceback.format_exc() - } - } diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py deleted file mode 100644 index e674354..0000000 --- a/laborious/utils/repository/opc_repository.py +++ /dev/null @@ -1,207 +0,0 @@ -from pathlib import Path -from asyncua.sync import Client -from asyncua.crypto.security_policies import SecurityPolicyBasic256 -from asyncua.ua import DataValue, Variant, VariantType -from logging import Logger -from datetime import datetime -from sientia_do.notifications.handlers import NotificationHandler -from sientia_do.notifications.models import NotificationLevel -import traceback - -data_type_map = { - 'float': { - 'converter': float, - 'opc_type': VariantType.Float, - }, - 'double': { - 'converter': float, - 'opc_type': VariantType.Double, - }, - 'int': { - 'converter': int, - 'opc_type': VariantType.Int32, - }, - 'bool': { - 'converter': bool, - 'opc_type': VariantType.Boolean, - }, - 'str': { - 'converter': str, - 'opc_type': VariantType.String, - } -} - - -class OpcRepository(): - def __init__(self, name: str, url: str, logger: Logger, notification_handler: NotificationHandler, - reconnection_interval: int = 60, server_uri: str = None, cert_path: str = None, - private_key_path: str = None, server_cert_path: str = None): - self.url = url - self.name = name - self.server_uri = server_uri - self.cert_path = cert_path - self.private_key_path = private_key_path - self.server_cert_path = server_cert_path - self.logger = logger - self.error_count = 0 - self.reconnection_interval = reconnection_interval - self.last_reconnection_time = None - self.notification_handler = notification_handler - self.client = None - - def set_security(self): - """ - Configures the security settings for the OPC UA client. - This method sets up the security policy, certificates, and timeouts - required for establishing a secure connection with the OPC UA server. - Raises: - ValueError: If either the certificate path or private key path is not provided. - Attributes: - cert_path (str): Path to the client's certificate file. - private_key_path (str): Path to the client's private key file. - server_cert_path (str, optional): Path to the server's certificate file. - server_uri (str): The URI of the server to be used as the application URI. - client (opcua.Client): The OPC UA client instance. - logger (logging.Logger): Logger instance for logging information. - Security Settings: - - Security Policy: Basic256 - - Secure Channel Timeout: 10,000,000 ms - - Session Timeout: 10,000,000 ms - """ - - if not all([self.cert_path, self.private_key_path]): - raise ValueError( - "Certificate and private key paths must be provided for secure connection.") - cert = Path(self.cert_path) - private_key = Path(self.private_key_path) - server_cert = Path( - self.server_cert_path) if self.server_cert_path else None - - self.client.application_uri = self.server_uri - self.logger.info('Setting security...') - self.client.set_security( - SecurityPolicyBasic256, - certificate=str(cert), - private_key=str(private_key), - server_certificate=str(server_cert) - ) - self.client.secure_channel_timeout = 10000000 - self.client.session_timeout = 10000000 - - def connect(self): - """ - Establishes a connection to the OPC server. - This method initializes the OPC client using the provided URL and - sets up security if a certificate path is specified. It then - attempts to connect to the server and logs the connection status. - Raises: - Exception: If the connection to the OPC server fails. - """ - - self.client = Client(self.url) - if self.cert_path: - self.set_security() - self.logger.info('Starting connection...') - return self.try_connect() - - def try_connect(self): - try: - self.last_reconnection_time = datetime.now() - self.client.connect() - return True - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id=f"OPC_CONNECTION_ERROR_{self.name}", - message=f"Failed to connect to OPC server: {e}", - block="opc_repository", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - self.logger.error(trace) - return False - - def disconnect(self): - if self.client is None: - return - self.client.disconnect() - self.client = None - self.logger.info('Disconnected from OPC server') - - def __del__(self): - try: - self.disconnect() - except Exception as e: - self.logger.error(f"Error in destructor: {e}") - - def validate_connection(self): - if self.client is None: - return self.connect() - - if self.error_count > 5: - self.logger.warning( - f"OPC server {self.name} will be disconnected due to multiple errors") - try: - self.disconnect() - except Exception as e: - trace = traceback.format_exc() - self.logger.error(f"Failed to disconnect from OPC server: {e}") - self.logger.error(trace) - self.logger.info( - f"Attempting to reconnect to OPC server {self.name}...") - return self.connect() - - if hasattr(self.client, 'aio_obj') and self.client.aio_obj.uaclient.protocol is None or \ - (hasattr(self.client.aio_obj.uaclient, 'protocol') and - self.client.aio_obj.uaclient.protocol.state == "closed"): - - self.logger.error( - f"OPC server {self.name} is not connected") - if (datetime.now() - self.last_reconnection_time).total_seconds( - ) > self.reconnection_interval: - self.logger.error( - f"Trying to reconnect to OPC server {self.name}...") - return self.try_connect() - - return False - - return True - - def write_data(self, node, value, data_type): - if not self.validate_connection(): - return - try: - node = self.client.get_node(node) - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id=f"OPC_WRITE_GET_NODE_ERROR_{self.name}", - message=f"Failed to get node from OPC server: {e}", - block="opc_repository", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - self.logger.error(trace) - self.error_count += 1 - return - - data = data_type_map[data_type]['converter'](value) - self.logger.info(f'Writing {data} - {type(data)} to {node}') - ua_data = DataValue( - Variant(data, data_type_map[data_type]['opc_type'])) - - try: - node.write_value(ua_data) - except Exception as e: - trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( - notification_id=f"OPC_WRITE_DATA_ERROR_{self.name}", - message=f"Failed to write data to OPC server: {e}", - block="opc_repository", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - self.logger.error(trace) - self.error_count += 1 - return - self.error_count = 0 diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py deleted file mode 100644 index 425f59f..0000000 --- a/laborious/workflows/predictions_batch.py +++ /dev/null @@ -1,89 +0,0 @@ -from temporalio import workflow - -with workflow.unsafe.imports_passed_through(): - from laborious.activities.activities import Activities - from typing import Any - from laborious.utils.policies import retry_policy - from datetime import timedelta - - -@workflow.defn(name="predictions_batch") -class PredictionsBatch(): - @workflow.run - async def run(self, input_data: dict[str, Any]): - """ - This workflow runs a batch of predictions based on the input data. - - The workflow executes in two main steps: - 1. Prepares the activity with schedule and model information - 2. Loads data using a custom query and executes the prediction process - - Args: - input_data (dict[str, Any]): The input data for the workflow. - Contains the following keys: - schedule_name (str): The name of the schedule. - model_name (str): The name of the model. - model_id (int): The id of the model. - query (str): The SQL query to be executed to load data. - schema (dict, optional): The schema definition for the data. - table_name (str, optional): The name of the table to process. - input_filters (dict, optional): Filters to be applied during prediction. - mlflow_transform_filters (dict, optional): Filters to be applied during prediction. - mlflow_predict_filters (dict, optional): Filters to be applied during prediction. - model_retention (int, optional): The model retention period in minutes. - path_priority (list[str]): The path priority. - Returns: - None - - Raises: - Exception: If any of the required parameters are missing or if the workflow fails. - """ - - await workflow.execute_local_activity_method( - Activities.prepare_activity, - { - 'schedule_name': input_data['schedule_name'], - 'model_name': input_data['model_name'], - 'model_id': input_data['model_id'], - 'workflow_name': 'predictions_batch' - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) - ) - - data = await workflow.execute_local_activity_method( - Activities.load_custom_query, - input_data['query'], - retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) - ) - - # Prepare input for prediction_process workflow - prediction_input = { - 'data': data, - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'model_id': input_data['model_id'], - 'model_name': input_data['model_name'], - 'input_filters': input_data.get('input_filters', { - 'EMPTY_DATA': { - 'POLICY': 'STOP' - } - }), - 'mlflow_transform_filters': input_data.get('mlflow_transform_filters', { - 'API_ERROR': { - 'POLICY': 'STOP' - } - }), - 'mlflow_predict_filters': input_data.get('mlflow_predict_filters', { - 'API_ERROR': { - 'POLICY': 'STOP' - } - }), - 'model_retention': input_data.get('model_retention', 60), - 'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), - 'opc_output_config': input_data.get('opc_output_config', {}) - } - - await workflow.execute_child_workflow( - 'prediction_process', prediction_input) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py deleted file mode 100644 index 3b1fad7..0000000 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ /dev/null @@ -1,95 +0,0 @@ -from temporalio import workflow - -with workflow.unsafe.imports_passed_through(): - from laborious.activities.activities import Activities - from typing import Any - from datetime import timedelta - from laborious.utils.policies import retry_policy - - -@workflow.defn(name="format_and_export_prediction") -class FormatAndExportPrediction(): - @workflow.run - async def run(self, input_data: dict[str, Any]): - """ - This workflow formats and exports predictions based on path_flag: - - If path_flag is None: formats prediction - using input data, timestamp, model_id and confidence - - If path_flag exists: creates default prediction - with timestamp, model_id, confidence and comment - Finally exports formatted prediction to postgres table - Args: - input_data(dict[str, Any]): The input data for the workflow. - Contains the following keys: - path_flag(str): The path flag to determine the type of prediction to format - data(dict[str, Any]): The data to format - prediction_confidence(float): The prediction confidence to be registered - timestamp(str): The timestamp of the prediction, synchronized with the data - model_id(int): The model id of the prediction - model_name(str): The model name of the prediction - model_retention(str): The model retention of the prediction - comment(str): The comment to be registered - schema(str): The schema of the prediction - table_name(str): The table name of the prediction - opc_output_config(dict[str, Any]): The opc output config of the prediction - - Returns: - bool: True if the workflow was successful, False otherwise. - """ - path_flag = input_data['path_flag'] - data = input_data['data'] - prediction_confidence = input_data['prediction_confidence'] - - if path_flag is None: - # proceed with formatting and exporting - prediction = await workflow.execute_local_activity_method( - Activities.format_prediction, - { - 'data': data, - 'timestamp': input_data['timestamp'], - 'model_id': input_data['model_id'], - 'prediction_confidence': prediction_confidence, - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) - ) - - else: - # create default prediction - prediction = await workflow.execute_local_activity_method( - Activities.format_default_prediction, - { - 'timestamp': input_data['timestamp'], - 'model_id': input_data['model_id'], - 'prediction_confidence': prediction_confidence, - 'comment': input_data['comment'] - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) - ) - - # write to postgres - postgres_holder = workflow.execute_activity_method( - Activities.export_data_to_postgres, - { - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'data': prediction - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) - ) - - # write to opc - opc_holder = workflow.execute_activity_method( - Activities.write_opc_data, - { - 'opc_output_config': input_data['opc_output_config'], - 'data': prediction - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) - ) - - await postgres_holder - await opc_holder diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py deleted file mode 100644 index 3164639..0000000 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ /dev/null @@ -1,233 +0,0 @@ -from temporalio import workflow - -with workflow.unsafe.imports_passed_through(): - from laborious.activities.activities import Activities - from typing import Any - from laborious.utils.policies import retry_policy - from datetime import timedelta - - -@workflow.defn(name="prediction_process") -class PredictionProcess(): - @workflow.run - async def run(self, input_data: dict[str, Any]): - """ - This workflow runs a prediction process based on the input data. - - The workflow executes in two main steps: - 1. Prepares the activity with schedule and model information - 2. Loads data using a custom query and executes the prediction process - - Args: - input_data (dict[str, Any]): The input data for the workflow. - Contains the following keys: - data (dict[str, Any]): The data to be used for the prediction. - schema (str): The schema of the table. - table_name (str): The name of the table. - model_id (int): The id of the model. - input_filters (dict, optional): Filters to be applied during prediction. - mlflow_transform_filters (dict, optional): Filters to be applied during prediction. - mlflow_predict_filters (dict, optional): Filters to be applied during prediction. - model_name (str): The name of the model. - model_retention (int, optional): The model retention period in minutes. - path_priority (list[str]): The path priority. - opc_output_config (dict[str, Any]): The opc output config of the prediction. - Returns: - None - - Raises: - Exception: If any of the required parameters are missing or if the workflow fails. - """ - - data = input_data['data'] - model_id = input_data['model_id'] - model_name = input_data['model_name'] - model_retention = input_data['model_retention'] - - last_timestamp = await workflow.execute_local_activity_method( - Activities.get_last_timestamp, - { - 'data': data - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) - - path_flag, confidence, comment = await workflow.execute_local_activity_method( - Activities.input_gate, - { - 'filters': input_data['input_filters'], - 'data': data, - 'path_priority': input_data['path_priority'] - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) - - if await self.path_flag_handler( - data, path_flag, input_data, confidence, last_timestamp, comment - ): - return - - response_data = await workflow.execute_local_activity_method( - Activities.request_transform, - { - 'data': data, - 'model_name': model_name, - 'model_retention': model_retention - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) - - path_flag, confidence, comment = await workflow.execute_local_activity_method( - Activities.mlflow_response_gate, - { - 'filters': input_data['mlflow_transform_filters'], - 'data': response_data, - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) - - if await self.path_flag_handler( - data, path_flag, input_data, confidence, last_timestamp, comment - ): - return - - transformed_data = response_data['content'] - - path_flag, confidence, comment = await workflow.execute_local_activity_method( - Activities.mlflow_content_gate, - { - 'filters': input_data['mlflow_transform_filters'], - 'data': transformed_data, - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) - - if await self.path_flag_handler( - data, path_flag, input_data, confidence, last_timestamp, comment - ): - return - - response_data = await workflow.execute_local_activity_method( - Activities.request_predict, - { - 'data': transformed_data, - 'model_name': model_name, - 'model_retention': model_retention - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) - - path_flag, confidence, comment = await workflow.execute_local_activity_method( - Activities.mlflow_response_gate, - { - 'filters': input_data['mlflow_predict_filters'], - 'data': response_data, - 'type': 'predict', - 'path_priority': input_data['path_priority'] - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) - - if await self.path_flag_handler( - data, path_flag, input_data, confidence, last_timestamp, comment - ): - return - - await workflow.execute_child_workflow( - 'format_and_export_prediction', - { - 'path_flag': path_flag, - 'data': response_data['content'], - 'prediction_confidence': confidence, - 'timestamp': last_timestamp, - 'model_id': model_id, - 'model_name': model_name, - 'model_retention': model_retention, - 'opc_output_config': input_data['opc_output_config'], - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'comment': comment - } - ) - - async def path_flag_handler(self, data: dict[str, Any], path_flag: str, - input_data: dict[str, Any], confidence: int, - last_timestamp: str, comment: str): - """ - This function handles the path flag and the confidence of the prediction. - It returns True if the prediction should be stopped. If path_flag is 'repeat', - it repeats the last prediction. - If path_flag is 'continue', it calls the write workflow. If path_flag is 'stop', - it stops the prediction process. - Args: - data (dict[str, Any]): The data to be used for the prediction. - path_flag (str): The path flag to determine the type of prediction to format - confidence (int): The confidence of the prediction - schema (str): The schema of the prediction - table_name (str): The table name of the prediction - model_id (int): The model id of the prediction - last_timestamp (str): The timestamp of the last prediction - model_name (str): The model name of the prediction - model_retention (int): The model retention of the prediction - comment (str): The comment of the prediction - Returns: - bool: True if the prediction should be stopped, False otherwise. - """ - - schema = input_data['schema'] - table_name = input_data['table_name'] - model_id = input_data['model_id'] - model_name = input_data['model_name'] - model_retention = input_data['model_retention'] - - path_flag = path_flag.upper() if path_flag else None - - if path_flag == 'STOP': - return True - - elif path_flag == 'REPEAT': - # repeat last prediction - await workflow.execute_activity_method( - Activities.repeat_last_prediction, - { - 'schema': schema, - 'table_name': table_name, - 'model_id': model_id - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) - return True - - elif path_flag == 'CONTINUE': - # call write workflow - await workflow.execute_child_workflow( - 'format_and_export_prediction', - { - 'path_flag': path_flag, - 'data': data, - 'prediction_confidence': confidence, - 'timestamp': last_timestamp, - 'model_id': model_id, - 'model_name': model_name, - 'model_retention': model_retention, - 'schema': schema, - 'table_name': table_name, - 'comment': comment, - 'opc_output_config': input_data['opc_output_config'] - } - ) - return True - - return False diff --git a/laborious/__init__.py b/orchestrator/__init__.py similarity index 100% rename from laborious/__init__.py rename to orchestrator/__init__.py diff --git a/laborious/activities/__init__.py b/orchestrator/activities/__init__.py similarity index 100% rename from laborious/activities/__init__.py rename to orchestrator/activities/__init__.py diff --git a/orchestrator/activities/activities.py b/orchestrator/activities/activities.py new file mode 100644 index 0000000..9579ec3 --- /dev/null +++ b/orchestrator/activities/activities.py @@ -0,0 +1,47 @@ +from temporalio import activity, workflow +from temporalio.client import Client + +with workflow.unsafe.imports_passed_through(): + from orchestrator.activities.couchbase import Couchbase + from orchestrator.activities.temporal_manager import TemporalManager + from orchestrator.activities.slot_manager import SlotManager + from typing import Any + from logging import Logger + from sientia_do.notifications.handlers import NotificationHandler + + +class Activities(Couchbase, TemporalManager, SlotManager): + + def __init__(self, + temporal_client: Client, + couchbase_config: dict[str, Any], + redis_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler): + + # Initialize parent classes + Couchbase.__init__(self, connection_string=couchbase_config['connection_string'], + username=couchbase_config['username'], + password=couchbase_config['password'], + logger=logger, + notification_handler=notification_handler) + + TemporalManager.__init__(self, + temporal_client=temporal_client, + logger=logger, + notification_handler=notification_handler) + + SlotManager.__init__(self, + host=redis_config['host'], + port=redis_config['port'], + username=redis_config['username'], + password=redis_config['password'], + logger=logger, + notification_handler=notification_handler) + + @activity.defn(name="prepare_activity") + async def prepare_activity(self, input_data: dict[str, Any]): + await super().prepare_activity(input_data) + + def shutdown(self): + Couchbase.shutdown(self) diff --git a/orchestrator/activities/couchbase.py b/orchestrator/activities/couchbase.py new file mode 100644 index 0000000..4aeec5b --- /dev/null +++ b/orchestrator/activities/couchbase.py @@ -0,0 +1,94 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + from typing import Any + from logging import Logger + from datetime import timedelta + import json + import traceback + from couchbase.auth import PasswordAuthenticator + from couchbase.cluster import Cluster + from couchbase.options import ClusterOptions + from sientia_do.notifications.handlers import NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.temporal.activities.base import BaseActivity + + +class Couchbase(BaseActivity): + def __init__(self, connection_string: str, username: str, + password: str, logger: Logger, + notification_handler: NotificationHandler): + + self.connection_string = connection_string + self.username = username + self.password = password + + logger.info("Initializing Couchbase connection...") + self.cluster = Cluster( + connection_string, + ClusterOptions( + authenticator=PasswordAuthenticator( + username=username, + password=password + ) + ) + ) + + logger.info("Awaiting Couchbase connection...") + self.cluster.wait_until_ready(timeout=timedelta(seconds=10)) + + logger.info("Couchbase connection ready") + + BaseActivity.__init__(self, + logger=logger, + notification_handler=notification_handler) + + def shutdown(self): + try: + self.cluster.close() + except Exception as e: + self.logger.error("Failed to close Couchbase connection: %s", e) + + def __del__(self): + self.shutdown() + + @activity.defn(name="load_query_from_couchbase") + async def load_query_from_couchbase(self, input_data: dict[str, Any]) -> list[dict[str, Any]]: + """ + Load a query from couchbase + + Args: + input_data (dict[str, Any]): The input data containing the query to execute + + Returns: + list[dict[str, Any]]: The result of the query + """ + query = input_data['query'] + + self.logger.info("Executing couchbase query: %s", query) + + try: + result = self.cluster.query(query) + except Exception as e: + trace = traceback.format_exc() + self.notification_handler.build_and_send_notification( + notification_id="COUCHBASE_LOAD_QUERY_ERROR", + message=f"Failed to execute couchbase query: {e}", + block="load_query_from_couchbase", + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + + self.logger.error(trace) + raise e + + rows = [] + + for row in result.rows(): + rows.append(row) + + self.logger.info("Fetched %d rows from couchbase", len(rows)) + self.logger.debug("Rows: \n %s", + json.dumps(rows, indent=4, sort_keys=True)) + + return rows diff --git a/orchestrator/activities/formatters.py b/orchestrator/activities/formatters.py new file mode 100644 index 0000000..4317834 --- /dev/null +++ b/orchestrator/activities/formatters.py @@ -0,0 +1,130 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + from sientia_do.temporal.activities.base import BaseActivity + from typing import Any + from logging import Logger + from sientia_do.notifications.handlers import NotificationHandler + + +class Formatters(BaseActivity): + def __init__(self, logger: Logger, notification_handler: NotificationHandler): + BaseActivity.__init__(self, logger=logger, + notification_handler=notification_handler) + + @activity.defn(name="process_schedules") + async def process_schedules(self, input_data: dict[str, Any]): + pipelines = input_data['pipelines'] + + schedule_config = {} + + for pipeline in pipelines: + if pipeline['workflow_type'] == 'scouter': + schedule_config[pipeline['schedule_name']] = scouter(pipeline) + + return schedule_config + + +def common_config(config: dict[str, Any]): + return { + "workflow_type": "scouter", + "schedule_name": config['schedule_name'], + "frequency": config.get('frequency', '1m'), + "max_retry_policy": config.get('max_retry_policy', 1), + + "model_id": config['model_id'], + "model_name": config['model_name'], + } + + +def scouter(config: dict[str, Any]): + filters = {} + for f in config['filters']: + filters[f['filter_name']] = { + "policy": f['policy'] + } + + tags = {} + for tag in config['read_tags']: + tags[tag['tag_name']] = { + "aggr_func": tag.get('aggr_func', 'lts'), + "data_range": tag.get('data_range', [-100, 100]) + } + + return { + **common_config(config), + + "topic": f"raw_{config['schedule_name']}", + "trigger_laborious": False, + "filters": filters, + "schema": "sientia_data", + "table_name": "laborious_data", + "retention_time": config.get('tag_retention_minutes', 60) * 60, + "model_tags": tags + } + + +def overlap_filter_config(base_filter_config: dict[str, Any], config: dict[str, Any]): + for fil in config['filters']: + base_filter_config[fil['filter_name']] = { + "policy": fil['policy'], + "config": fil.get('config', {}) + } + + return base_filter_config + + +def predictions_batch(config: dict[str, Any]): + tags = {} + for tag in config['write_tags']: + if tag['server_name'] not in tags: + tags[tag['server_name']] = {} + + tag_type = tag['type'] + + if tag_type == 'prediction' or tag_type == 'confidence': + tag_type_str = f"{tag_type}_tags" + + if tag_type_str not in tags[tag['server_name']]: + tags[tag['server_name']][tag_type_str] = {} + + tags[tag['server_name']][tag_type_str][tag['addr']] = { + "data_type": tag.get('data_type', 'float'), + } + + path_priority = config.get('path_priority', ["STOP", "CONTINUE", "REPEAT"]) + + for priority in path_priority[:]: + if priority not in ["STOP", "CONTINUE", "REPEAT"]: + path_priority.remove(priority) + + if len(path_priority) != 3: + for priority in ["STOP", "CONTINUE", "REPEAT"]: + if priority not in path_priority: + path_priority.append(priority) + + return { + **common_config(config), + + "query": config['query'], + "schema": "sientia_data", + "table_name": "predictions", + "retention_time": config.get('model_retention_minutes', 60) * 60, + "opc_output_config": tags, + "input_filters": overlap_filter_config({ + "EMPTY_DATA": { + "policy": "STOP" + } + }, config['input_filters']), + "mlflow_transform_filters": overlap_filter_config({ + "API_ERROR": { + "policy": "STOP" + } + }, config['mlflow_transform_filters']), + "mlflow_predict_filters": overlap_filter_config({ + "API_ERROR": { + "policy": "STOP" + } + }, config['mlflow_predict_filters']), + "path_priority": path_priority + } diff --git a/laborious/utils/__init__.py b/orchestrator/activities/notification.py similarity index 100% rename from laborious/utils/__init__.py rename to orchestrator/activities/notification.py diff --git a/orchestrator/activities/slot_manager.py b/orchestrator/activities/slot_manager.py new file mode 100644 index 0000000..f6ab941 --- /dev/null +++ b/orchestrator/activities/slot_manager.py @@ -0,0 +1,73 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + from typing import Any + import json + from logging import Logger + from sientia_do.temporal.activities.redis_base import Redis + from sientia_do.notifications.handlers import NotificationHandler + + +class SlotManager(Redis): + + def __init__(self, host: str, port: int, + username: str, password: str, + logger: Logger, notification_handler: NotificationHandler): + + Redis.__init__(self, host, port, username, + password, logger, notification_handler) + + @activity.defn(name="load_opc_slots") + async def load_opc_slots(self) -> dict[str, Any]: + """ + Load all OPC slots from Redis + + Returns: + dict[str, Any]: A dictionary of OPC slots + """ + + self.logger.info("Loading OPC slots...") + + opc_slots = {} + + slot_keys = self.redis_client.keys("slot:opc_tags:*") + + if slot_keys: + decoded_keys = [key.decode('utf-8') for key in slot_keys] + values = self.redis_client.mget(decoded_keys) + + for i, key in enumerate(decoded_keys): + value = values[i] + if value is not None: + try: + opc_slots[key] = value.decode('utf-8') + except (UnicodeDecodeError, AttributeError): + opc_slots[key] = value + else: + opc_slots[key] = None + + self.logger.info(f"Loaded {len(opc_slots)} OPC slots") + + self.logger.debug("OPC slots: \n %s", + json.dumps(opc_slots, indent=4, sort_keys=True)) + + return opc_slots + + @activity.defn(name="load_active_ingestors") + async def load_active_ingestors(self) -> list[str]: + """ + Load all active ingestors from Redis + + Returns: + list[str]: A list of active ingestors + """ + + self.logger.info("Loading active ingestors...") + + active_ingestors = self.redis_client.keys("heartbeat:ingestor:*") + + self.logger.info(f"Loaded {len(active_ingestors)} active ingestors") + + self.logger.debug("Active ingestors: \n %s", active_ingestors) + + return [ingestor.decode('utf-8') for ingestor in active_ingestors] diff --git a/orchestrator/activities/temporal_manager.py b/orchestrator/activities/temporal_manager.py new file mode 100644 index 0000000..daa8fc5 --- /dev/null +++ b/orchestrator/activities/temporal_manager.py @@ -0,0 +1,66 @@ +from temporalio import activity, workflow +from temporalio.client import Client + +with workflow.unsafe.imports_passed_through(): + from typing import Any + from logging import Logger + from sientia_do.notifications.handlers import NotificationHandler + from sientia_do.temporal.activities.base import BaseActivity + from google.protobuf.json_format import MessageToDict + import base64 + import json + + +class TemporalManager(BaseActivity): + def __init__(self, temporal_client: Client, logger: Logger, + notification_handler: NotificationHandler): + + self.temporal_client = temporal_client + + BaseActivity.__init__(self, + logger=logger, + notification_handler=notification_handler) + + @activity.defn(name="load_schedule") + async def load_schedule(self) -> dict[str, Any]: + """ + Load all orchestrated schedules from Temporal. Filters by search attribute + "Orchestrated" set to "true" and returns a dictionary of schedule_id: + {frequency, data, handle} + + Returns: + dict[str, Any]: A dictionary of orchestrated schedules + """ + + self.logger.info("Getting orchestrated schedules...") + + orchestrated_schedules = {} + + async for schedule in await self.temporal_client.list_schedules(): + search_attrs = getattr(schedule, "search_attributes", {}) + if search_attrs.get("Orchestrated", ["false"]) == ["true"]: + schedule_id = schedule.id + + handle = self.temporal_client.get_schedule(schedule_id) + + desc = await handle.describe() + + for arg in desc.schedule.action.args: + data = MessageToDict(arg)['data'] + data = base64.b64decode(data).decode('utf-8') + + frequency = desc.schedule.spec.intervals[0].every.seconds + + orchestrated_schedules[schedule_id] = { + 'frequency': frequency, + 'data': json.loads(data), + 'handle': handle + } + + self.logger.info("Found %d orchestrated schedules", + len(orchestrated_schedules)) + + self.logger.debug("Orchestrated schedules: %s", + orchestrated_schedules) + + return orchestrated_schedules diff --git a/laborious/utils/filters/__init__.py b/orchestrator/utils/__init__.py similarity index 100% rename from laborious/utils/filters/__init__.py rename to orchestrator/utils/__init__.py diff --git a/laborious/utils/connectors_config.py b/orchestrator/utils/connectors_config.py similarity index 100% rename from laborious/utils/connectors_config.py rename to orchestrator/utils/connectors_config.py diff --git a/laborious/worker/__init__.py b/orchestrator/worker/__init__.py similarity index 100% rename from laborious/worker/__init__.py rename to orchestrator/worker/__init__.py diff --git a/laborious/worker/worker.py b/orchestrator/worker/worker.py similarity index 100% rename from laborious/worker/worker.py rename to orchestrator/worker/worker.py diff --git a/laborious/workflows/__init__.py b/orchestrator/workflows/__init__.py similarity index 100% rename from laborious/workflows/__init__.py rename to orchestrator/workflows/__init__.py diff --git a/orchestrator/workflows/orchestrator.py b/orchestrator/workflows/orchestrator.py new file mode 100644 index 0000000..ca64d87 --- /dev/null +++ b/orchestrator/workflows/orchestrator.py @@ -0,0 +1,69 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from orchestrator.activities.activities import Activities + from typing import Any + from datetime import timedelta + from sientia_do.temporal.utils.policies import retry_policy + + +@workflow.defn(name="orchestrator") +class Orchestrator: + @workflow.run + async def run(self, input_data: dict[str, Any]): + + input_data['workflow_name'] = 'orchestrator' + + await workflow.execute_local_activity_method( + Activities.prepare_activity, + { + 'workflow_name': input_data['workflow_name'], + 'schedule_name': input_data['schedule_name'], + 'model_name': '-', + 'model_id': '-' + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60) + ) + + pipeline_config_handler = workflow.execute_local_activity_method( + Activities.load_query_from_couchbase, + { + 'query': input_data['pipelines_query'] + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60) + ) + + opc_servers_handler = workflow.execute_local_activity_method( + Activities.load_query_from_couchbase, + { + 'query': input_data['opc_servers_query'] + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60) + ) + + orchestrated_schedules_handler = workflow.execute_local_activity_method( + Activities.load_schedule, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60) + ) + + slot_config_handler = workflow.execute_local_activity_method( + Activities.load_opc_slots, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60) + ) + + active_ingestors_handler = workflow.execute_local_activity_method( + Activities.load_active_ingestors, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60) + ) + + pipeline_config = await pipeline_config_handler + orchestrated_schedules = await orchestrated_schedules_handler + slot_config = await slot_config_handler + opc_servers = await opc_servers_handler + active_ingestors = await active_ingestors_handler diff --git a/requirements.txt b/requirements.txt index de4ff2a..655d2f0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ temporalio psycopg2-binary sqlalchemy redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git +couchbase +/home/grezewave/Documents/projects/sientia/sientia-dataops-library/ diff --git a/samples.json b/samples.json new file mode 100644 index 0000000..1b66dd5 --- /dev/null +++ b/samples.json @@ -0,0 +1,35 @@ +{ + "models": { + "1": { + "name": "Demo Model-Demo2" + } + }, + "pipelines": { + "1": { + "name": "scouter-opcua-pipeline", + "model_id": 1, + "workflow_type": "scouter", + "frequency": "5s", + "max_retry_policy": 1, + + "read_tags": [ + { + "tag_name": "Counter", + "aggr_func": "avg", + "data_range": [-100, 100] + } + ], + "filters": [ + { + "filter_name": "OUT_OF_BOUNDS_FILTER", + "policy": "DISCARD" + }, + { + "filter_name": "NULL_VALUES_FILTER", + "policy": "DISCARD" + } + ], + "tag_retention_minutes": 60 + } + } +} \ No newline at end of file diff --git a/simulator/Dockerfile b/simulator/Dockerfile deleted file mode 100644 index d467676..0000000 --- a/simulator/Dockerfile +++ /dev/null @@ -1,30 +0,0 @@ -# syntax=docker/dockerfile:1.4 - -FROM python:3.11-slim - -# Enable use of SSH agent/socket -# This line enables SSH during build -# (don't forget the syntax header above) -RUN apt-get update && apt-get install -y git openssh-client && rm -rf /var/lib/apt/lists/* - -# Use build-time SSH mount for Git clone -# The SSH key will NOT remain in the image -# IMPORTANT: this block requires BuildKit -# and the --ssh flag during docker build - -# SSH config to skip host key check (safe in CI/local dev) -RUN mkdir -p /root/.ssh && echo "StrictHostKeyChecking no" > /root/.ssh/config - -WORKDIR /app - -# Clone using SSH -ARG GIT_REPO -ARG GIT_BRANCH=main - -# Mount SSH key just for this RUN -RUN --mount=type=ssh git clone --branch ${GIT_BRANCH} ${GIT_REPO} . - -# Install requirements if exists -RUN if [ -f requirements.txt ]; then pip install --no-cache-dir -r requirements.txt; fi - -CMD ["python", "server.py"] diff --git a/test.ipynb b/test.ipynb new file mode 100644 index 0000000..8a0ebc6 --- /dev/null +++ b/test.ipynb @@ -0,0 +1,876 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 4, + "id": "a287fa45", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Initializing Couchbase connection...\n", + "Awaiting Couchbase connection...\n", + "Couchbase connection ready\n" + ] + } + ], + "source": [ + "from orchestrator.activities.couchbase import Couchbase\n", + "from unittest.mock import MagicMock\n", + "logger = MagicMock(info=MagicMock(side_effect=print), debug=MagicMock(side_effect=print))\n", + "couchbase = Couchbase(\n", + " connection_string=\"couchbase://localhost\",\n", + " username=\"sientia\",\n", + " password=\"sientia\",\n", + " logger=logger,\n", + " notification_handler=MagicMock()\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "d9a0f9c4", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Executing couchbase query: %s \n", + "SELECT\n", + " pipelines.*,\n", + " models as model\n", + "FROM\n", + " `pipelines`\n", + "JOIN\n", + " `models` ON KEYS pipelines.model_id;\n", + "\n", + "Fetched %d rows from couchbase 1\n", + "Rows: \n", + " %s [\n", + " {\n", + " \"filters\": [\n", + " {\n", + " \"filter_name\": \"OUT_OF_BOUNDS_FILTER\",\n", + " \"policy\": \"DISCARD\"\n", + " },\n", + " {\n", + " \"filter_name\": \"NULL_VALUES_FILTER\",\n", + " \"policy\": \"DISCARD\"\n", + " }\n", + " ],\n", + " \"frequency\": \"5s\",\n", + " \"max_retry_policy\": 1,\n", + " \"model\": {\n", + " \"name\": \"Demo Model-Demo2\"\n", + " },\n", + " \"model_id\": \"1\",\n", + " \"name\": \"scouter-opcua-pipeline\",\n", + " \"read_tags\": [\n", + " {\n", + " \"aggr_func\": \"avg\",\n", + " \"data_range\": [\n", + " -100,\n", + " 100\n", + " ],\n", + " \"tag_name\": \"Counter\"\n", + " }\n", + " ],\n", + " \"tag_retention_minutes\": 60,\n", + " \"workflow_type\": \"scouter\"\n", + " }\n", + "]\n", + "[{'filters': [{'filter_name': 'OUT_OF_BOUNDS_FILTER', 'policy': 'DISCARD'}, {'filter_name': 'NULL_VALUES_FILTER', 'policy': 'DISCARD'}], 'frequency': '5s', 'max_retry_policy': 1, 'model': {'name': 'Demo Model-Demo2'}, 'model_id': '1', 'name': 'scouter-opcua-pipeline', 'read_tags': [{'aggr_func': 'avg', 'data_range': [-100, 100], 'tag_name': 'Counter'}], 'tag_retention_minutes': 60, 'workflow_type': 'scouter'}]\n" + ] + } + ], + "source": [ + "query = \"\"\"\n", + "SELECT\n", + " pipelines.*,\n", + " models as model\n", + "FROM\n", + " `pipelines`\n", + "JOIN\n", + " `models` ON KEYS pipelines.model_id;\n", + "\"\"\"\n", + "if __name__ == \"__main__\":\n", + "\n", + "\n", + " result = await couchbase.load_query_from_couchbase({\n", + " \"query\": query\n", + " })\n", + "\n", + " print(result)\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "7d01f160", + "metadata": {}, + "outputs": [], + "source": [ + "from temporalio import client\n", + "from orchestrator.activities.temporal_manager import TemporalManager\n", + "import os\n", + "from unittest.mock import MagicMock\n", + "\n", + "host = \"localhost:7233\"\n", + "logger = MagicMock(info=MagicMock(side_effect=print), debug=MagicMock(side_effect=print))\n", + "\n", + "temporal_client = await client.Client.connect(\n", + " target_host=host,\n", + " namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')\n", + ")\n", + "\n", + "manager = TemporalManager(temporal_client=temporal_client,\n", + " logger=logger,\n", + " notification_handler=MagicMock())\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": 81, + "id": "bb750ae6", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 81, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import asyncio\n", + "from datetime import timedelta\n", + "from temporalio.client import (\n", + " Client,\n", + " Schedule,\n", + " ScheduleActionStartWorkflow,\n", + " ScheduleIntervalSpec,\n", + " ScheduleSpec,\n", + ")\n", + "from temporalio.common import TypedSearchAttributes, SearchAttributeKey, SearchAttributePair\n", + "\n", + "\n", + "customer_id_key = SearchAttributeKey.for_keyword(\"Orchestrated\")\n", + "search_attributes = TypedSearchAttributes([\n", + " SearchAttributePair(customer_id_key, \"true\")\n", + "])\n", + "await temporal_client.create_schedule(\n", + " \"meu-schedule-id5\",\n", + " Schedule(\n", + " action=ScheduleActionStartWorkflow(\n", + " 'scouter-test2',\n", + " {\n", + " 'args': {\n", + " 'arg1': 'value1'\n", + " }\n", + " },\n", + " id=\"workflow-id-unico\",\n", + " task_queue=\"nome-da-task-queue\",\n", + " ),\n", + " spec=ScheduleSpec(\n", + " intervals=[ScheduleIntervalSpec(every=timedelta(minutes=10))]\n", + " )\n", + " ),\n", + " search_attributes=search_attributes,\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 86, + "id": "1bd82225", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Getting orchestrated schedules...\n", + "Schedule: %s ScheduleListDescription(id='meu-schedule-id5', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test2'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=, _value_type=), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=, payload_codec=None, failure_converter_class=, payload_converter=, failure_converter=), raw_entry=schedule_id: \"meu-schedule-id5\"\n", + "search_attributes {\n", + " indexed_fields {\n", + " key: \"Orchestrated\"\n", + " value {\n", + " metadata {\n", + " key: \"type\"\n", + " value: \"Text\"\n", + " }\n", + " metadata {\n", + " key: \"encoding\"\n", + " value: \"json/plain\"\n", + " }\n", + " data: \"\\\"true\\\"\"\n", + " }\n", + " }\n", + "}\n", + "info {\n", + " spec {\n", + " interval {\n", + " interval {\n", + " seconds: 600\n", + " }\n", + " }\n", + " }\n", + " workflow_type {\n", + " name: \"scouter-test2\"\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548800\n", + " }\n", + " future_action_times {\n", + " seconds: 1748549400\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550000\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550600\n", + " }\n", + " future_action_times {\n", + " seconds: 1748551200\n", + " }\n", + "}\n", + ")\n", + "Search attributes: %s {'Orchestrated': ['true']}\n", + "Schedule: %s ScheduleListDescription(id='meu-schedule-id4', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test2'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=, _value_type=), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=, payload_codec=None, failure_converter_class=, payload_converter=, failure_converter=), raw_entry=schedule_id: \"meu-schedule-id4\"\n", + "search_attributes {\n", + " indexed_fields {\n", + " key: \"Orchestrated\"\n", + " value {\n", + " metadata {\n", + " key: \"type\"\n", + " value: \"Text\"\n", + " }\n", + " metadata {\n", + " key: \"encoding\"\n", + " value: \"json/plain\"\n", + " }\n", + " data: \"\\\"true\\\"\"\n", + " }\n", + " }\n", + "}\n", + "info {\n", + " spec {\n", + " interval {\n", + " interval {\n", + " seconds: 600\n", + " }\n", + " }\n", + " }\n", + " workflow_type {\n", + " name: \"scouter-test2\"\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548800\n", + " }\n", + " future_action_times {\n", + " seconds: 1748549400\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550000\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550600\n", + " }\n", + " future_action_times {\n", + " seconds: 1748551200\n", + " }\n", + "}\n", + ")\n", + "Search attributes: %s {'Orchestrated': ['true']}\n", + "Schedule: %s ScheduleListDescription(id='meu-schedule-id3', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test2'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=, _value_type=), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=, payload_codec=None, failure_converter_class=, payload_converter=, failure_converter=), raw_entry=schedule_id: \"meu-schedule-id3\"\n", + "search_attributes {\n", + " indexed_fields {\n", + " key: \"Orchestrated\"\n", + " value {\n", + " metadata {\n", + " key: \"type\"\n", + " value: \"Text\"\n", + " }\n", + " metadata {\n", + " key: \"encoding\"\n", + " value: \"json/plain\"\n", + " }\n", + " data: \"\\\"true\\\"\"\n", + " }\n", + " }\n", + "}\n", + "info {\n", + " spec {\n", + " interval {\n", + " interval {\n", + " seconds: 600\n", + " }\n", + " }\n", + " }\n", + " workflow_type {\n", + " name: \"scouter-test2\"\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548800\n", + " }\n", + " future_action_times {\n", + " seconds: 1748549400\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550000\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550600\n", + " }\n", + " future_action_times {\n", + " seconds: 1748551200\n", + " }\n", + "}\n", + ")\n", + "Search attributes: %s {'Orchestrated': ['true']}\n", + "Schedule: %s ScheduleListDescription(id='meu-schedule-id2', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test2'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 50, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 50, 0, 37623, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='workflow-id-unico-2025-05-29T19:50:00Z', first_execution_run_id='01971d98-265a-7285-b46f-cf09bcf2d301'))], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=, _value_type=), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=, payload_codec=None, failure_converter_class=, payload_converter=, failure_converter=), raw_entry=schedule_id: \"meu-schedule-id2\"\n", + "search_attributes {\n", + " indexed_fields {\n", + " key: \"Orchestrated\"\n", + " value {\n", + " metadata {\n", + " key: \"type\"\n", + " value: \"Text\"\n", + " }\n", + " metadata {\n", + " key: \"encoding\"\n", + " value: \"json/plain\"\n", + " }\n", + " data: \"\\\"true\\\"\"\n", + " }\n", + " }\n", + "}\n", + "info {\n", + " spec {\n", + " interval {\n", + " interval {\n", + " seconds: 600\n", + " }\n", + " }\n", + " }\n", + " workflow_type {\n", + " name: \"scouter-test2\"\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548200\n", + " }\n", + " actual_time {\n", + " seconds: 1748548200\n", + " nanos: 37623285\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"workflow-id-unico-2025-05-29T19:50:00Z\"\n", + " run_id: \"01971d98-265a-7285-b46f-cf09bcf2d301\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548800\n", + " }\n", + " future_action_times {\n", + " seconds: 1748549400\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550000\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550600\n", + " }\n", + " future_action_times {\n", + " seconds: 1748551200\n", + " }\n", + "}\n", + ")\n", + "Search attributes: %s {'Orchestrated': ['true']}\n", + "Schedule: %s ScheduleListDescription(id='meu-schedule-id', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 0, 0, 37788, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='workflow-id-unico-2025-05-29T19:00:00Z', first_execution_run_id='01971d6a-5fa1-70f8-8371-60ad11587b77'))], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=, _value_type=), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=, payload_codec=None, failure_converter_class=, payload_converter=, failure_converter=), raw_entry=schedule_id: \"meu-schedule-id\"\n", + "search_attributes {\n", + " indexed_fields {\n", + " key: \"Orchestrated\"\n", + " value {\n", + " metadata {\n", + " key: \"type\"\n", + " value: \"Text\"\n", + " }\n", + " metadata {\n", + " key: \"encoding\"\n", + " value: \"json/plain\"\n", + " }\n", + " data: \"\\\"true\\\"\"\n", + " }\n", + " }\n", + "}\n", + "info {\n", + " spec {\n", + " interval {\n", + " interval {\n", + " seconds: 600\n", + " }\n", + " }\n", + " }\n", + " workflow_type {\n", + " name: \"scouter-test\"\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748545200\n", + " }\n", + " actual_time {\n", + " seconds: 1748545200\n", + " nanos: 37788631\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"workflow-id-unico-2025-05-29T19:00:00Z\"\n", + " run_id: \"01971d6a-5fa1-70f8-8371-60ad11587b77\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548800\n", + " }\n", + " future_action_times {\n", + " seconds: 1748549400\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550000\n", + " }\n", + " future_action_times {\n", + " seconds: 1748550600\n", + " }\n", + " future_action_times {\n", + " seconds: 1748551200\n", + " }\n", + "}\n", + ")\n", + "Search attributes: %s {'Orchestrated': ['true']}\n", + "Schedule: %s ScheduleListDescription(id='laborious_test', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='predictions_batch'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=60), offset=datetime.timedelta(0))], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 53, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 53, 0, 35750, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:53:00Z', first_execution_run_id='01971d9a-e57f-700b-953b-bdb3b05810e2')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 54, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 54, 0, 35780, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:54:00Z', first_execution_run_id='01971d9b-cfde-7f39-8b38-b7e62fee1f82')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 55, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 55, 0, 34329, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:55:00Z', first_execution_run_id='01971d9c-ba3c-7d27-9db7-72759ebabc76')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 56, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 56, 0, 49214, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:56:00Z', first_execution_run_id='01971d9d-a4a9-7b55-a77b-fd450e768ccf')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 57, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 57, 0, 36109, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:57:00Z', first_execution_run_id='01971d9e-8eff-7608-9b10-0ed1875df9fe'))], next_action_times=[datetime.datetime(2025, 5, 29, 19, 58, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 59, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 1, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 2, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[]), search_attributes={}, data_converter=DataConverter(payload_converter_class=, payload_codec=None, failure_converter_class=, payload_converter=, failure_converter=), raw_entry=schedule_id: \"laborious_test\"\n", + "info {\n", + " spec {\n", + " interval {\n", + " interval {\n", + " seconds: 60\n", + " }\n", + " phase {\n", + " }\n", + " }\n", + " }\n", + " workflow_type {\n", + " name: \"predictions_batch\"\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548380\n", + " }\n", + " actual_time {\n", + " seconds: 1748548380\n", + " nanos: 35750962\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"laborious_test-2025-05-29T19:53:00Z\"\n", + " run_id: \"01971d9a-e57f-700b-953b-bdb3b05810e2\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548440\n", + " }\n", + " actual_time {\n", + " seconds: 1748548440\n", + " nanos: 35780414\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"laborious_test-2025-05-29T19:54:00Z\"\n", + " run_id: \"01971d9b-cfde-7f39-8b38-b7e62fee1f82\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548500\n", + " }\n", + " actual_time {\n", + " seconds: 1748548500\n", + " nanos: 34329717\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"laborious_test-2025-05-29T19:55:00Z\"\n", + " run_id: \"01971d9c-ba3c-7d27-9db7-72759ebabc76\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548560\n", + " }\n", + " actual_time {\n", + " seconds: 1748548560\n", + " nanos: 49214684\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"laborious_test-2025-05-29T19:56:00Z\"\n", + " run_id: \"01971d9d-a4a9-7b55-a77b-fd450e768ccf\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548620\n", + " }\n", + " actual_time {\n", + " seconds: 1748548620\n", + " nanos: 36109875\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"laborious_test-2025-05-29T19:57:00Z\"\n", + " run_id: \"01971d9e-8eff-7608-9b10-0ed1875df9fe\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548680\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548740\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548800\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548860\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548920\n", + " }\n", + "}\n", + ")\n", + "Search attributes: %s {}\n", + "Schedule: %s ScheduleListDescription(id='scouter-opcua-pipeline', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=30), offset=datetime.timedelta(0))], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 55, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 55, 0, 26130, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:55:00Z', first_execution_run_id='01971d9c-ba35-7b2a-b596-effe9939c7c7')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 55, 30, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 55, 30, 26550, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:55:30Z', first_execution_run_id='01971d9d-2f66-7012-bd77-e5809b8c4c7a')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 56, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 56, 0, 39848, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:56:00Z', first_execution_run_id='01971d9d-a4a1-7ba9-9205-21f157fe6e54')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 56, 30, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 56, 30, 39791, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:56:30Z', first_execution_run_id='01971d9e-19d2-7358-96d3-261e203faaa7')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 57, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 57, 0, 28221, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:57:00Z', first_execution_run_id='01971d9e-8ef8-72af-a903-1391b5e06c2f'))], next_action_times=[datetime.datetime(2025, 5, 29, 19, 57, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 58, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 58, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 59, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 59, 30, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[]), search_attributes={}, data_converter=DataConverter(payload_converter_class=, payload_codec=None, failure_converter_class=, payload_converter=, failure_converter=), raw_entry=schedule_id: \"scouter-opcua-pipeline\"\n", + "info {\n", + " spec {\n", + " interval {\n", + " interval {\n", + " seconds: 30\n", + " }\n", + " phase {\n", + " }\n", + " }\n", + " }\n", + " workflow_type {\n", + " name: \"scouter\"\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548500\n", + " }\n", + " actual_time {\n", + " seconds: 1748548500\n", + " nanos: 26130839\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:55:00Z\"\n", + " run_id: \"01971d9c-ba35-7b2a-b596-effe9939c7c7\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548530\n", + " }\n", + " actual_time {\n", + " seconds: 1748548530\n", + " nanos: 26550164\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:55:30Z\"\n", + " run_id: \"01971d9d-2f66-7012-bd77-e5809b8c4c7a\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548560\n", + " }\n", + " actual_time {\n", + " seconds: 1748548560\n", + " nanos: 39848794\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:56:00Z\"\n", + " run_id: \"01971d9d-a4a1-7ba9-9205-21f157fe6e54\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548590\n", + " }\n", + " actual_time {\n", + " seconds: 1748548590\n", + " nanos: 39791709\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:56:30Z\"\n", + " run_id: \"01971d9e-19d2-7358-96d3-261e203faaa7\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1748548620\n", + " }\n", + " actual_time {\n", + " seconds: 1748548620\n", + " nanos: 28221068\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:57:00Z\"\n", + " run_id: \"01971d9e-8ef8-72af-a903-1391b5e06c2f\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548650\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548680\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548710\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548740\n", + " }\n", + " future_action_times {\n", + " seconds: 1748548770\n", + " }\n", + "}\n", + ")\n", + "Search attributes: %s {}\n", + "Found %d orchestrated schedules 5\n" + ] + } + ], + "source": [ + "schedules = await manager.load_schedule({})" + ] + }, + { + "cell_type": "code", + "execution_count": 87, + "id": "a4c777dd", + "metadata": {}, + "outputs": [], + "source": [ + "from google.protobuf.json_format import MessageToDict\n", + "import base64\n", + "import json\n", + "\n", + "\n", + "schedules_config = {}\n", + "for schedule in schedules:\n", + " id = schedule.id\n", + "\n", + " handle = temporal_client.get_schedule_handle(id)\n", + "\n", + " desc = await handle.describe()\n", + "\n", + " for arg in desc.schedule.action.args:\n", + " data = MessageToDict(arg)['data']\n", + " data = base64.b64decode(data).decode('utf-8')\n", + "\n", + " frequency = desc.schedule.spec.intervals[0].every.seconds\n", + "\n", + " schedules_config[id] = {\n", + " 'frequency': frequency,\n", + " 'data': json.loads(data),\n", + " 'handle': handle\n", + " }\n", + " \n", + " \n", + " \n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": 88, + "id": "988d1718", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'meu-schedule-id5': {'frequency': 600,\n", + " 'data': {'args': {'arg1': 'value1'}},\n", + " 'handle': },\n", + " 'meu-schedule-id4': {'frequency': 600,\n", + " 'data': {},\n", + " 'handle': },\n", + " 'meu-schedule-id3': {'frequency': 600,\n", + " 'data': {},\n", + " 'handle': },\n", + " 'meu-schedule-id2': {'frequency': 600,\n", + " 'data': {},\n", + " 'handle': },\n", + " 'meu-schedule-id': {'frequency': 600,\n", + " 'data': {},\n", + " 'handle': }}" + ] + }, + "execution_count": 88, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "schedules_config" + ] + }, + { + "cell_type": "code", + "execution_count": 85, + "id": "c12f5e75", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "meu-schedule-id5 {'frequency': 600, 'data': {'args': {'arg1': 'value1'}}, 'handle': }\n", + "meu-schedule-id4 {'frequency': 1200, 'data': {'args': {'arg1': 'value1'}}, 'handle': }\n", + "meu-schedule-id4\n", + "meu-schedule-id3 {'frequency': 600, 'data': {}, 'handle': }\n", + "meu-schedule-id2 {'frequency': 600, 'data': {}, 'handle': }\n", + "meu-schedule-id {'frequency': 600, 'data': {}, 'handle': }\n" + ] + } + ], + "source": [ + "from temporalio.client import Client, ScheduleUpdateInput, ScheduleUpdate, ScheduleSpec\n", + "\n", + "\n", + "for id, schedule in schedules_config.items():\n", + " print(id, schedule)\n", + " \n", + " if schedule['frequency'] != 600:\n", + " print(id)\n", + "\n", + " handle = schedule['handle']\n", + " \n", + " async def update_schedule(input: ScheduleUpdateInput) -> ScheduleUpdate:\n", + " schedule_action = input.description.schedule.action\n", + " \n", + " if hasattr(schedule_action, 'args'):\n", + " schedule_action.args = [{}]\n", + " \n", + " # Atualiza o intervalo de execução\n", + " input.description.schedule.spec.intervals = [\n", + " ScheduleIntervalSpec(every=timedelta(minutes=10))\n", + " ]\n", + " \n", + " return ScheduleUpdate(schedule=input.description.schedule)\n", + "\n", + " await handle.update(update_schedule)\n", + " \n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "87eb10c8", + "metadata": {}, + "outputs": [], + "source": [ + "from redis import Redis\n", + "\n", + "redis = Redis(host='localhost', port=6379)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "fe4ad9dc", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Chave: slot:opc_tags:2, Valor: {\n", + "\"slot2\": \"value\"\n", + "}\n", + "Chave: slot:opc_tags:1, Valor: {\n", + "\"slot1\": \"value\"\n", + "}\n" + ] + } + ], + "source": [ + "matching_keys = redis.keys(\"slot:opc_tags:*\")\n", + "\n", + "if matching_keys:\n", + " decoded_keys = [key.decode('utf-8') for key in matching_keys]\n", + " values = redis.mget(decoded_keys)\n", + "\n", + " items = {}\n", + " for i, key in enumerate(decoded_keys):\n", + " value = values[i]\n", + " if value is not None:\n", + " try:\n", + " items[key] = value.decode('utf-8')\n", + " except (UnicodeDecodeError, AttributeError):\n", + " items[key] = value\n", + " else:\n", + " items[key] = None\n", + "\n", + " for key, value in items.items():\n", + " print(f\"Chave: {key}, Valor: {value}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "87a9dc89", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/laborious/activities/__init__.py b/tests/laborious/activities/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py deleted file mode 100644 index 3b2ef49..0000000 --- a/tests/laborious/activities/test_activities.py +++ /dev/null @@ -1,193 +0,0 @@ -from pytest import mark -from unittest.mock import patch, MagicMock, ANY -from laborious.activities.activities import Activities -from laborious.activities.postgres import Postgres -from laborious.activities.mlflow import MLFlow -from laborious.activities.gates import Gates -from laborious.activities.opc import OPC - - -@patch('laborious.activities.activities.Postgres.__init__') -@patch('laborious.activities.activities.MLFlow.__init__') -@patch('laborious.activities.activities.OPC.__init__') -@patch('laborious.activities.activities.Gates.__init__') -def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init): - - postgres_config = { - 'host': 'localhost', - 'port': 5432, - 'user': 'postgres', - 'password': 'postgres', - 'dbname': 'postgres', - 'min_connections': 1, - 'max_connections': 10 - } - - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' - } - - opc_config = { - 'bootstrap_servers': 'localhost:9092', - 'polling_time': 1000, - 'group_id': 'test-group' - } - - logger = MagicMock() - notification_handler = MagicMock() - - activities = Activities( - postgres_config=postgres_config, - mlflow_config=mlflow_config, - opc_config=opc_config, - logger=logger, - notification_handler=notification_handler - ) - - assert isinstance(activities, Activities) - assert isinstance(activities, Postgres) - assert isinstance(activities, MLFlow) - assert isinstance(activities, OPC) - assert isinstance(activities, Gates) - - mock_postgres_init.assert_called_once_with( - ANY, - host=postgres_config['host'], - port=postgres_config['port'], - user=postgres_config['user'], - password=postgres_config['password'], - dbname=postgres_config['dbname'], - min_connections=postgres_config['min_connections'], - max_connections=postgres_config['max_connections'], - logger=logger, - notification_handler=notification_handler - ) - - mock_mlflow_init.assert_called_once_with( - ANY, - mlflow_host=mlflow_config['host'], - mlflow_port=mlflow_config['port'], - mlflow_username=mlflow_config['username'], - mlflow_password=mlflow_config['password'], - logger=logger, - notification_handler=notification_handler - ) - - mock_opc_init.assert_called_once_with( - ANY, - opc_servers=opc_config, - logger=logger, - notification_handler=notification_handler - ) - - mock_gates_init.assert_called_once_with( - ANY, - logger=logger, - notification_handler=notification_handler - ) - - -@mark.asyncio -@patch('laborious.activities.activities.Postgres.__init__') -@patch('laborious.activities.activities.MLFlow.__init__') -@patch('laborious.activities.activities.OPC.__init__') -async def test_prepare_activity(_mock_opc_init, - _mock_mlflow_init, _mock_postgres_init): - postgres_config = { - 'host': 'localhost', - 'port': 5432, - 'user': 'postgres', - 'password': 'postgres', - 'dbname': 'postgres', - 'min_connections': 1, - 'max_connections': 10 - } - - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' - } - - opc_config = { - 'bootstrap_servers': 'localhost:9092', - 'polling_time': 1000, - 'group_id': 'test-group' - } - - logger = MagicMock() - notification_handler = MagicMock() - - activities = Activities( - postgres_config=postgres_config, - mlflow_config=mlflow_config, - opc_config=opc_config, - logger=logger, - notification_handler=notification_handler - ) - - input_data = { - 'workflow_name': 'test-workflow-name', - 'schedule_name': 'test-schedule-name', - 'model_name': 'test-model-name', - 'model_id': 'test-model-id' - } - - await activities.prepare_activity(input_data) - - assert activities.notification_handler.base_notification.pipeline_name == input_data[ - 'workflow_name'] - assert activities.notification_handler.base_notification.schedule_name == input_data[ - 'schedule_name'] - assert activities.notification_handler.base_notification.model_name == input_data[ - 'model_name'] - assert activities.notification_handler.base_notification.model_id == input_data[ - 'model_id'] - - -@patch('laborious.activities.activities.Postgres', return_value=MagicMock()) -@patch('laborious.activities.activities.MLFlow', return_value=MagicMock()) -@patch('laborious.activities.activities.OPC', return_value=MagicMock()) -def test_shutdown(mock_opc_init, - _mock_mlflow_init, mock_postgres_init): - postgres_config = { - 'host': 'localhost', - 'port': 5432, - 'user': 'postgres', - 'password': 'postgres', - 'dbname': 'postgres', - 'min_connections': 1, - 'max_connections': 10 - } - - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' - } - - opc_config = { - 'bootstrap_servers': 'localhost:9092', - 'polling_time': 1000, - 'group_id': 'test-group' - } - - logger = MagicMock() - notification_handler = MagicMock() - - activities = Activities( - postgres_config=postgres_config, - mlflow_config=mlflow_config, - opc_config=opc_config, - logger=logger, - notification_handler=notification_handler - ) - - activities.shutdown() - mock_opc_init.shutdown.assert_called_once() - mock_postgres_init.close.assert_called_once() diff --git a/tests/laborious/activities/test_base.py b/tests/laborious/activities/test_base.py deleted file mode 100644 index 6978acb..0000000 --- a/tests/laborious/activities/test_base.py +++ /dev/null @@ -1,35 +0,0 @@ -from unittest.mock import MagicMock -from laborious.activities.base import BaseActivity -from pytest import fixture, mark -from sientia_do.notifications.models import Notification - - -@fixture -def base_activity(): - return BaseActivity( - logger=MagicMock(), - notification_handler=MagicMock(), - ) - - -@mark.asyncio -async def test_prepare_activity(base_activity): - base_activity.notification_handler.base_notification = Notification( - project="project", - pipeline="pipeline", - trigger="-", - model_name="-", - model_id="-", - ) - - await base_activity.prepare_activity({ - 'workflow_name': 'test_workflow', - '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" - assert base_activity.notification_handler.base_notification.pipeline_name == "test_workflow" diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py deleted file mode 100644 index 6b61c81..0000000 --- a/tests/laborious/activities/test_gates.py +++ /dev/null @@ -1,369 +0,0 @@ -from unittest.mock import MagicMock, ANY, patch -from pytest import fixture, mark -from sientia_do.notifications.models import NotificationLevel -from laborious.activities.gates import Gates - - -@fixture -def gates_activity(): - return Gates( - logger=MagicMock(), - notification_handler=MagicMock(), - ) - - -@mark.asyncio -async def test_input_gate_invalid_filter(gates_activity): - # Arrange - input_data = { - 'filters': { - 'INVALID_FILTER': {'POLICY': 'STOP'} - }, - 'data': {'value': [1, 2, 3]}, - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.input_gate(input_data) - - # Assert - assert result == (None, 0, "") - gates_activity.logger.error.assert_called_once_with( - "Filter INVALID_FILTER not found" - ) - - -@mark.asyncio -@patch('laborious.activities.gates.input_filter_functions') -async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity): - # Arrange - mock_input_filter_functions.__contains__.return_value = True - mock_input_filter_functions.__getitem__.return_value = MagicMock( - side_effect=Exception("Test error")) - input_data = { - 'filters': { - 'EMPTY_DATA': {'POLICY': 'STOP'} - }, - 'data': {'value': []}, - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.input_gate(input_data) - - # Assert - assert result == (None, 0, "") - gates_activity.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id="INTPUT_GATE_ERROR__EMPTY_DATA", - message="Error in filter EMPTY_DATA:{'POLICY': 'STOP'}: \n Test error", - block="input_gate", - level=NotificationLevel.ERROR, - attachment_content=ANY - ) - - -@mark.asyncio -async def test_input_gate_no_filters(gates_activity): - # Arrange - input_data = { - 'filters': {}, - 'data': {'value': [1, 2, 3]}, - 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'] - } - - # Act - result = await gates_activity.input_gate(input_data) - - # Assert - assert result == (None, 0, "") - gates_activity.logger.debug.assert_called() - - -@mark.asyncio -async def test_input_gate_with_filter(gates_activity): - # Arrange - input_data = { - 'filters': { - 'EMPTY_DATA': {'POLICY': 'STOP'} - }, - 'data': {'value': []}, - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.input_gate(input_data) - - # Assert - assert result == ('STOP', -1, "Input data with bad quality") - gates_activity.logger.debug.assert_called() - - -@mark.asyncio -async def test_mlflow_response_gate_invalid_filter(gates_activity): - # Arrange - input_data = { - 'filters': { - 'INVALID_FILTER': {'POLICY': 'STOP'} - }, - 'data': {'content': {'message': 'success'}}, - 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.mlflow_response_gate(input_data) - - # Assert - assert result == (None, 0, "") - - -@mark.asyncio -@patch('laborious.activities.gates.mlflow_response_filter_functions') -async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions, - gates_activity): - # Arrange - mock_mlflow_response_filter_functions.__contains__.return_value = True - mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock( - side_effect=Exception("Test error")) - input_data = { - 'filters': { - 'INVALID_FILTER': {'POLICY': 'STOP'} - }, - 'data': {'content': {'message': 'success'}}, - 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.mlflow_response_gate(input_data) - - # Assert - assert result == (None, 0, "") - gates_activity.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER", - message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error", - block="mlflow_gate", - level=NotificationLevel.ERROR, - attachment_content=ANY - ) - - -@mark.asyncio -async def test_mlflow_response_gate_no_filters(gates_activity): - # Arrange - input_data = { - 'filters': {}, - 'data': {'content': {'message': 'success'}}, - 'type': 'test', - 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'] - } - - # Act - result = await gates_activity.mlflow_response_gate(input_data) - - # Assert - assert result == (None, 0, "") - gates_activity.logger.debug.assert_called() - - -@mark.asyncio -async def test_mlflow_response_gate_with_filter(gates_activity): - # Arrange - input_data = { - 'filters': { - 'API_ERROR': {'POLICY': 'STOP'} - }, - 'data': { - 'success': False, - 'content': { - 'message': 'API error occurred', - 'traceback': 'error trace' - } - }, - 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.mlflow_response_gate(input_data) - - # Assert - assert result == ('STOP', -1, "API error occurred") - gates_activity.logger.debug.assert_called() - gates_activity.notification_handler.build_and_send_notification.assert_called() - - -@mark.asyncio -async def test_mlflow_content_gate_invalid_filter(gates_activity): - # Arrange - input_data = { - 'filters': { - 'INVALID_FILTER': {'POLICY': 'STOP'} - }, - 'data': {'value': [1, 2, 3]}, - 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.mlflow_content_gate(input_data) - - # Assert - assert result == (None, 0, "") - - -@mark.asyncio -@patch('laborious.activities.gates.mlflow_content_filter_functions') -async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions, - gates_activity): - # Arrange - mock_mlflow_content_filter_functions.__contains__.return_value = True - mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock( - side_effect=Exception("Test error")) - input_data = { - 'filters': { - 'API_ERROR': {'POLICY': 'STOP'} - }, - 'data': { - 'success': False, - 'content': { - 'message': 'API error occurred', - 'traceback': 'error trace' - } - }, - 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.mlflow_content_gate(input_data) - - # Assert - assert result == (None, 0, "") - gates_activity.logger.debug.assert_called() - gates_activity.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR", - message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error", - block="mlflow_gate", - level=NotificationLevel.ERROR, - attachment_content=ANY - ) - - -@mark.asyncio -async def test_mlflow_content_gate_no_filters(gates_activity): - # Arrange - input_data = { - 'filters': {}, - 'data': {'value': [1, 2, 3]}, - 'type': 'test', - 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'] - } - - # Act - result = await gates_activity.mlflow_content_gate(input_data) - - # Assert - assert result == (None, 0, "") - gates_activity.logger.debug.assert_called() - - -@mark.asyncio -async def test_mlflow_content_gate_with_filter(gates_activity): - # Arrange - input_data = { - 'filters': { - 'NAN_VALUES': {'POLICY': 'STOP'} - }, - 'data': {'value': [None, None, None]}, - 'type': 'test', - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'] - } - - # Act - result = await gates_activity.mlflow_content_gate(input_data) - - # Assert - assert result == ( - 'STOP', -1, "Transformed data not passed the content filter") - gates_activity.logger.debug.assert_called() - gates_activity.notification_handler.build_and_send_notification.assert_called() - - -@mark.asyncio -async def test_format_prediction(gates_activity): - # Arrange - input_data = { - 'data': {'prediction': [1], 'response_time': [0.1]}, - 'timestamp': '2023-05-26 11:12:27', - 'model_id': 'test_model', - 'prediction_confidence': 0.9 - } - - # Act - result = await gates_activity.format_prediction(input_data) - - # Assert - assert result['prediction'] == {0: 1} - assert result['response_time'] == {0: ANY} - assert result['timestamp'] == {0: '2023-05-26 11:12:27'} - assert result['model_id'] == {0: 'test_model'} - assert result['prediction_confidence'] == {0: 0.9} - assert result['prediction_status'] == {0: 'Good'} - assert result['comments'] == {0: ""} - gates_activity.logger.debug.assert_called() - - -@mark.asyncio -async def test_format_default_prediction(gates_activity): - # Arrange - input_data = { - 'timestamp': '2023-05-26 11:12:27', - 'model_id': 'test_model', - 'prediction_confidence': 0.1, - 'comment': 'Test comment' - } - - # Act - result = await gates_activity.format_default_prediction(input_data) - - # Assert - assert result['prediction'] == {0: 0} - assert result['response_time'] == {0: 0} - assert result['timestamp'] == {0: '2023-05-26 11:12:27'} - assert result['model_id'] == {0: 'test_model'} - assert result['prediction_confidence'] == {0: 0.1} - assert result['prediction_status'] == {0: 'Bad'} - assert result['comments'] == {0: 'Test comment'} - gates_activity.logger.debug.assert_called() - - -@mark.asyncio -async def test_get_last_timestamp_with_data(gates_activity): - # Arrange - input_data = { - 'data': { - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'] - } - } - - # Act - result = await gates_activity.get_last_timestamp(input_data) - - # Assert - assert result == '2023-05-26 11:12:28' - - -@mark.asyncio -async def test_get_last_timestamp_no_data(gates_activity): - # Arrange - input_data = { - 'data': {} - } - - # Act - result = await gates_activity.get_last_timestamp(input_data) - - # Assert - assert isinstance(result, str) # Should be a timestamp string - assert len(result) > 0 diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py deleted file mode 100644 index 5834cb4..0000000 --- a/tests/laborious/activities/test_mlflow.py +++ /dev/null @@ -1,120 +0,0 @@ -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 = 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 - - # 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 deleted file mode 100644 index d012778..0000000 --- a/tests/laborious/activities/test_opc.py +++ /dev/null @@ -1,196 +0,0 @@ -from unittest.mock import patch, MagicMock, ANY, call -from pytest import fixture, mark -from laborious.activities.opc import NotificationLevel - -from laborious.activities.opc import OPC - - -@patch("laborious.activities.opc.OpcRepository") -def test___init__(mock_opc_repository): - mock_logger = MagicMock() - server1 = MagicMock() - server2 = MagicMock() - mock_opc_repository.side_effect = [server1, server2] - mock_notification_handler = MagicMock() - servers = { - 'server1': { - 'url': 'http://localhost:8080', - 'server_uri': 'opc.tcp://localhost:4840', - 'cert_path': '', - 'private_key_path': '', - 'server_cert_path': '', - 'reconnection_interval': 60, - }, - 'server2': { - 'url': 'http://localhost:8080', - 'server_uri': 'opc.tcp://localhost:4840', - 'cert_path': '', - 'private_key_path': '', - 'server_cert_path': '', - 'reconnection_interval': 60, - } - } - opc = OPC( - opc_servers=servers, - logger=mock_logger, - notification_handler=mock_notification_handler - ) - - assert opc.opc_servers == servers - assert opc.logger == mock_logger - assert opc.notification_handler == mock_notification_handler - assert opc.opc_repository['server1'] == server1 - assert opc.opc_repository['server2'] == server2 - - mock_opc_repository.assert_has_calls([ - call( - name="server1", - url="http://localhost:8080", - logger=mock_logger, - server_uri="opc.tcp://localhost:4840", - cert_path="", - private_key_path="", - server_cert_path="", - notification_handler=mock_notification_handler, - reconnection_interval=60, - ), - ]) - mock_opc_repository.assert_has_calls([ - call( - name="server2", - url="http://localhost:8080", - logger=mock_logger, - server_uri="opc.tcp://localhost:4840", - cert_path="", - private_key_path="", - server_cert_path="", - notification_handler=mock_notification_handler, - reconnection_interval=60, - ) - ]) - - server1.connect.assert_called_once() - server2.connect.assert_called_once() - - -@fixture -@patch("laborious.activities.opc.OpcRepository") -def opc(_mock_opc_repository): - servers = { - 'server1': { - 'url': 'http://localhost:8080', - 'server_uri': 'opc.tcp://localhost:4840', - 'cert_path': '', - 'private_key_path': '', - 'server_cert_path': '', - 'reconnection_interval': 60, - } - } - return OPC( - opc_servers=servers, - logger=MagicMock(), - notification_handler=MagicMock() - ) - - -WRITE_DATA_CASES = [ - ('tag1', 'int', 50), - ('tag2', 'float', 50.5), - ('tag3', 'bool', True), - ('tag4', 'string', 'test'), -] - - -@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES) -def test_write_data_success(opc, tag, data_type, data): - opc.write_data(server='server1', tag=tag, data=data, - data_type=data_type, tag_type='prediction') - opc.opc_repository['server1'].write_data.assert_called_once_with( - tag, data, data_type) - - -def test_write_data_exception(opc): - opc.opc_repository['server1'].write_data.side_effect = Exception( - "Test error") - opc.write_data(server='server1', tag='tag1', data=50, - data_type='int', tag_type='prediction') - opc.notification_handler.build_and_send_notification.assert_called_once_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_success(opc): - # Arrange - input_data = { - 'data': { - 'prediction': [0.75], - 'prediction_confidence': [0.95] - }, - 'opc_output_config': { - 'server1': { - 'prediction_tags': { - 'tag1': {'data_type': 'float'} - }, - 'confidence_tags': { - 'tag2': {'data_type': 'float'} - } - } - } - } - - # Act - opc.write_data = MagicMock() - await opc.write_opc_data(input_data) - - # Assert - opc.write_data.assert_has_calls([ - call( - server='server1', - tag='tag1', - data=0.75, - data_type='float', - tag_type='prediction' - )]) - opc.write_data.assert_has_calls([ - call( - server='server1', - tag='tag2', - data=0.95, - data_type='float', - tag_type='confidence' - ) - ]) - assert opc.write_data.call_count == 2 - - -@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['server1'].write_data.assert_not_called() - - -def test_shutdown(opc): - opc.shutdown() - opc.opc_repository['server1'].disconnect.assert_called_once() diff --git a/tests/laborious/activities/test_postgres.py b/tests/laborious/activities/test_postgres.py deleted file mode 100644 index e4a4545..0000000 --- a/tests/laborious/activities/test_postgres.py +++ /dev/null @@ -1,159 +0,0 @@ -from unittest.mock import MagicMock, patch -from pytest import fixture, mark -import pandas as pd -from laborious.activities.postgres import Postgres - - -@fixture -@patch("laborious.activities.postgres.create_engine") -def postgres_activity(_mock_create_engine): - return Postgres( - host="localhost", - port=5432, - user="test_user", - password="test_password", - dbname="test_db", - min_connections=1, - max_connections=5, - logger=MagicMock(), - notification_handler=MagicMock() - ) - - -@mark.asyncio -@patch("laborious.activities.postgres.read_sql_query") -async def test_load_custom_query_none_data(mock_read_sql_query, postgres_activity): - query = "SELECT * FROM test_table LIMIT 1" - mock_read_sql_query.return_value = None - - result = await postgres_activity.load_custom_query(query) - - assert isinstance(result, dict) - assert len(result) == 0 - - -@mark.asyncio -@patch("laborious.activities.postgres.read_sql_query") -async def test_load_custom_query_date_converted(mock_read_sql_query, postgres_activity): - query = "SELECT * FROM test_table LIMIT 1" - mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]}) - mock_data['date'] = pd.to_datetime('2022-01-01') - - mock_read_sql_query.return_value = mock_data - - result = await postgres_activity.load_custom_query(query) - - assert isinstance(result, dict) - assert len(result) == 3 - assert "column1" in result - assert "column2" in result - assert "date" in result - assert result['date'] == {0: '2022-01-01 00:00:00'} - - -@mark.asyncio -@patch("laborious.activities.postgres.read_sql_query") -async def test_load_custom_query_success(mock_read_sql_query, postgres_activity): - query = "SELECT * FROM test_table LIMIT 1" - mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]}) - - mock_read_sql_query.return_value = mock_data - - result = await postgres_activity.load_custom_query(query) - - assert isinstance(result, dict) - assert len(result) == 2 - assert "column1" in result - assert "column2" in result - postgres_activity.logger.info.assert_called() - - -@mark.asyncio -async def test_load_custom_query_error(postgres_activity): - query = "SELECT * FROM non_existent_table" - error_msg = "Table not found" - - with patch("laborious.activities.postgres.read_sql_query", side_effect=ValueError(error_msg)): - result = await postgres_activity.load_custom_query(query) - - assert isinstance(result, dict) - assert len(result) == 0 - postgres_activity.notification_handler.build_and_send_notification.assert_called_once() - postgres_activity.logger.error.assert_called() - - -@mark.asyncio -async def test_repeat_last_prediction_success(postgres_activity): - query_items = { - "schema": "public", - "table_name": "predictions", - "model": 1 - } - - with patch("sqlalchemy.orm.session.Session.execute") as mock_execute: - await postgres_activity.repeat_last_prediction(query_items) - - mock_execute.assert_called_once() - postgres_activity.logger.info.assert_called() - - -@mark.asyncio -async def test_repeat_last_prediction_error(postgres_activity): - query_items = { - "schema": "public", - "table_name": "predictions", - "model": 1 - } - error_msg = "Database error" - - with patch("sqlalchemy.orm.session.Session.execute", side_effect=ValueError(error_msg)): - await postgres_activity.repeat_last_prediction(query_items) - - postgres_activity.notification_handler.build_and_send_notification.assert_called_once() - postgres_activity.logger.error.assert_called() - - -@mark.asyncio -async def test_export_data_to_postgres_success(postgres_activity): - input_data = { - "schema": "public", - "table_name": "test_table", - "data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]}) - } - - with patch("laborious.activities.postgres.DataFrame.to_sql") as mock_to_sql: - await postgres_activity.export_data_to_postgres(input_data) - - mock_to_sql.assert_called_once() - postgres_activity.logger.debug.assert_called() - - -@mark.asyncio -async def test_export_data_to_postgres_error(postgres_activity): - input_data = { - "schema": "public", - "table_name": "test_table", - "data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]}) - } - error_msg = "Export failed" - - with patch("laborious.activities.postgres.DataFrame.to_sql", side_effect=ValueError(error_msg)): - await postgres_activity.export_data_to_postgres(input_data) - - postgres_activity.notification_handler.build_and_send_notification.assert_called_once() - postgres_activity.logger.error.assert_called() - - -@mark.asyncio -async def test_close(postgres_activity): - postgres_activity.close() - - postgres_activity.engine.dispose.assert_called_once() - - -@mark.asyncio -async def test_del(postgres_activity): - postgres_activity.close = MagicMock() - postgres_activity.__del__() - - postgres_activity.close.assert_called_once() diff --git a/tests/laborious/utils/__init__.py b/tests/laborious/utils/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/laborious/utils/filters/__init__.py b/tests/laborious/utils/filters/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py deleted file mode 100644 index edcbcd6..0000000 --- a/tests/laborious/utils/filters/test_conditional_filters.py +++ /dev/null @@ -1,30 +0,0 @@ -from pandas import DataFrame - -from laborious.utils.filters.conditional_filters import ( - filter_specific_variables_null_values, - filter_empty_data -) - - -def test_filter_specific_variables_null_values(): - assert filter_specific_variables_null_values( - DataFrame( - {'variable': ['variable1', 'variable2'], 'value': [1, 2]}), - config={'VARIABLES': ['variable2']}) is False - - -def test_filter_specific_variables_null_values_with_null_values(): - assert filter_specific_variables_null_values( - DataFrame( - {'variable': ['variable1', 'variable2'], 'value': [1, None]}), - config={'VARIABLES': ['variable2']}) is True - - -def test_filter_empty_data(): - assert filter_empty_data(DataFrame(), {}) is True - - -def test_filter_empty_data_with_data(): - assert filter_empty_data( - DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), - {}) is False diff --git a/tests/laborious/utils/filters/test_mlflow_filters.py b/tests/laborious/utils/filters/test_mlflow_filters.py deleted file mode 100644 index f9c61e9..0000000 --- a/tests/laborious/utils/filters/test_mlflow_filters.py +++ /dev/null @@ -1,22 +0,0 @@ -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/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py deleted file mode 100644 index a675edb..0000000 --- a/tests/laborious/utils/repository/test_model_repository.py +++ /dev/null @@ -1,278 +0,0 @@ -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'] is True - assert output['content'] == {'prediction': { - 0: 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/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py deleted file mode 100644 index ae9dd89..0000000 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ /dev/null @@ -1,259 +0,0 @@ -from unittest.mock import Mock, patch, MagicMock, ANY, call -from asyncua.crypto.security_policies import SecurityPolicyBasic256 -from pytest import fixture -from laborious.utils.repository.opc_repository import OpcRepository -from sientia_do.notifications.models import NotificationLevel -from datetime import datetime - - -@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, - notification_handler=Mock(), - reconnection_interval=60, - 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.reconnection_interval == 60 - assert opc_repository.client is None - assert opc_repository.last_reconnection_time is None - assert opc_repository.error_count == 0 - - -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.try_connect = MagicMock() - opc_repository.connect() - - opc_repository.try_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.try_connect = MagicMock() - opc_repository.set_security = MagicMock() - opc_repository.connect() - - opc_repository.try_connect.assert_called_once() - opc_repository.set_security.assert_not_called() - assert opc_repository.client == mock_client - - -def test_try_connect_sucess(opc_repository): - opc_repository.last_reconnection_time = None - opc_repository.client = MagicMock() - opc_repository.try_connect() - opc_repository.client.connect.assert_called_once() - assert opc_repository.last_reconnection_time is not None - - -def test_try_connect_fail(opc_repository): - opc_repository.last_reconnection_time = None - opc_repository.client = MagicMock() - opc_repository.client.connect.side_effect = Exception("Test error") - - opc_repository.try_connect() - - opc_repository.client.connect.assert_called_once() - assert opc_repository.last_reconnection_time is not None - opc_repository.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.name}", - message="Failed to connect to OPC server: Test error", - block="opc_repository", - level=NotificationLevel.ERROR, - attachment_content=ANY - ) - - -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_validate_connection_none_client(opc_repository): - opc_repository.client = None - opc_repository.connect = MagicMock() - response = opc_repository.validate_connection() - assert response - opc_repository.connect.assert_called_once() - - -def test_validate_connection_error_count_disconnect_error(opc_repository): - opc_repository.error_count = 6 - opc_repository.client = MagicMock() - opc_repository.disconnect = MagicMock(side_effect=Exception("Test error")) - opc_repository.connect = MagicMock() - - response = opc_repository.validate_connection() - assert response == opc_repository.connect.return_value - opc_repository.disconnect.assert_called_once() - opc_repository.connect.assert_called_once() - opc_repository.logger.error.assert_has_calls( - [ - call("Failed to disconnect from OPC server: Test error"), - ] - ) - - -@patch('laborious.utils.repository.opc_repository.hasattr', return_value=True) -@patch('laborious.utils.repository.opc_repository.datetime', - MagicMock(now=MagicMock(return_value=datetime(2025, 1, 1, 0, 0, 0)))) -def test_validate_connection_lost_not_time_to_reconect(_mock_datetime, opc_repository): - opc_repository.error_count = 0 - opc_repository.client = MagicMock() - opc_repository.client.aio_obj.uaclient.protocol = None - opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0) - opc_repository.try_connect = MagicMock() - - response = opc_repository.validate_connection() - opc_repository.try_connect.assert_not_called() - assert response is False - - -@patch('laborious.utils.repository.opc_repository.hasattr', return_value=True) -@patch('laborious.utils.repository.opc_repository.datetime', - MagicMock(now=MagicMock(return_value=datetime(2025, 1, 1, 1, 0, 0)))) -def test_validate_connection_lost_time_to_reconect(_mock_datetime, opc_repository): - opc_repository.error_count = 0 - opc_repository.client = MagicMock() - opc_repository.client.aio_obj.uaclient.protocol = None - opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0) - opc_repository.try_connect = MagicMock() - - response = opc_repository.validate_connection() - opc_repository.try_connect.assert_called_once() - assert response == opc_repository.try_connect.return_value - - -def test_validate_connection_failed(opc_repository): - opc_repository.client = MagicMock() - opc_repository.error_count = 0 - - output = opc_repository.validate_connection() - assert output is True - - -def test_write_data_validate_connection_do_nothing(opc_repository): - opc_repository.validate_connection = MagicMock(return_value=True) - opc_repository.client = MagicMock() - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") - opc_repository.validate_connection.assert_called_once() - opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode") - - -def test_write_data_validate_connection_failed(opc_repository): - opc_repository.validate_connection = MagicMock(return_value=False) - opc_repository.client = MagicMock() - opc_repository.error_count = 0 - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") - opc_repository.validate_connection.assert_called_once() - opc_repository.client.get_node.assert_not_called() - - -def test_write_data_get_node_failed(opc_repository): - opc_repository.validate_connection = MagicMock(return_value=True) - opc_repository.client = MagicMock() - opc_repository.error_count = 0 - opc_repository.client.get_node.side_effect = Exception("Test error") - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") - opc_repository.validate_connection.assert_called_once() - opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode") - opc_repository.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.name}", - message="Failed to get node from OPC server: Test error", - block="opc_repository", - level=NotificationLevel.ERROR, - attachment_content=ANY - ) - assert opc_repository.error_count == 1 - - -def test_write_data(opc_repository, mock_client): - opc_repository.validate_connection = MagicMock(return_value=True) - 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_client.get_node.assert_called_once_with("ns=2;s=TestNode") - mock_node.write_value.assert_called_once() - opc_repository.logger.info.assert_called_once_with( - "Writing 42.0 - to " + str(mock_node)) - - -def test_write_data_write_value_failed(opc_repository, mock_client): - opc_repository.validate_connection = MagicMock(return_value=True) - opc_repository.client = mock_client - mock_node = MagicMock() - opc_repository.error_count = 0 - mock_client.get_node.return_value = mock_node - mock_node.write_value.side_effect = Exception("Test error") - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") - opc_repository.validate_connection.assert_called_once() - mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") - mock_node.write_value.assert_called_once() - opc_repository.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.name}", - message="Failed to write data to OPC server: Test error", - block="opc_repository", - level=NotificationLevel.ERROR, - attachment_content=ANY - ) - assert opc_repository.error_count == 1 diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py deleted file mode 100644 index b137bc2..0000000 --- a/tests/laborious/utils/test_connectors_config.py +++ /dev/null @@ -1,133 +0,0 @@ -from os import environ -from laborious.utils.connectors_config import (build_mlflow_config, - build_opc_config, - build_postgres_config) - - -def test_build_mlflow_config_with_env_vars(): - # Arrange - environ['MLFLOW_HOST'] = 'http://test-host' - environ['MLFLOW_PORT'] = '8080' - environ['MLFLOW_USERNAME'] = 'test-user' - environ['MLFLOW_PASSWORD'] = 'test-pass' - - # Act - config = build_mlflow_config() - - # Assert - assert config['host'] == 'http://test-host' - assert config['port'] == 8080 - assert config['username'] == 'test-user' - assert config['password'] == 'test-pass' - - -def test_build_mlflow_config_with_defaults(): - # Arrange - # Clear any existing env vars - environ.pop('MLFLOW_HOST', None) - environ.pop('MLFLOW_PORT', None) - environ.pop('MLFLOW_USERNAME', None) - environ.pop('MLFLOW_PASSWORD', None) - - # Act - config = build_mlflow_config() - - # Assert - assert config['host'] == 'http://localhost' - assert config['port'] == 5080 - assert config['username'] == 'aignosi' - assert config['password'] == 'aignosi' - - -def test_build_opc_config_with_env_vars(): - # Arrange - environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}' - - # Act - config = build_opc_config() - - # Assert - assert config['opc']['name'] == 'test-opc' - assert config['opc']['url'] == 'opc.tcp://test:4840' - - -def test_build_opc_config_with_individual_env_vars(): - # Arrange - environ.pop('OPC_CONFIG', None) - environ['OPC_NAME'] = 'test-name' - environ['OPC_URL'] = 'opc.tcp://test:4840' - environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840' - environ['OPC_RECONNECTION_INTERVAL'] = '300' - - # Act - config = build_opc_config() - - # Assert - assert config['opc']['name'] == 'test-name' - assert config['opc']['url'] == 'opc.tcp://test:4840' - assert config['opc']['server_uri'] == 'opc.tcp://test:4840' - assert config['opc']['reconnection_interval'] == 300 - - -def test_build_opc_config_with_defaults(): - # Arrange - environ.pop('OPC_CONFIG', None) - environ.pop('OPC_NAME', None) - environ.pop('OPC_URL', None) - environ.pop('OPC_SERVER_URI', None) - environ.pop('OPC_RECONNECTION_INTERVAL', None) - - # Act - config = build_opc_config() - - # Assert - assert config['opc']['name'] == 'opc' - assert config['opc']['url'] == 'opc.tcp://localhost:4840' - assert config['opc']['server_uri'] == 'opc.tcp://localhost:4840' - assert config['opc']['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 diff --git a/tests/laborious/utils/test_logger.py b/tests/laborious/utils/test_logger.py deleted file mode 100644 index cb68cb4..0000000 --- a/tests/laborious/utils/test_logger.py +++ /dev/null @@ -1,37 +0,0 @@ -import os -from unittest.mock import patch -import logging -import pytest -from laborious.utils.logger import get_logger - - -@pytest.fixture -def mock_env_vars(): - with patch.dict(os.environ, {}, clear=True): - yield - - -@pytest.mark.usefixtures("mock_env_vars") -@patch('laborious.utils.logger.logging.Formatter') -@patch('laborious.utils.logger.logging.StreamHandler') -def test_get_logger_defaults(mock_stream_handler, mock_formatter): - """Test logger creation with default settings""" - # Mock the StreamHandler and Formatter - - logger = get_logger('test_logger') - - # Verify logger settings - assert logger.name == 'test_logger' - assert logger.level == logging.INFO - - # Verify handler configuration - mock_stream_handler.return_value.setLevel.assert_called_once_with('INFO') - mock_stream_handler.return_value.setFormatter.assert_called_once() - - # Verify formatter configuration - mock_formatter.assert_called_once_with( - '%(asctime)s - %(name)s - %(levelname)s - %(message)s' - ) - - # Verify handler was added to logger - assert len(logger.handlers) == 1 diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py deleted file mode 100644 index f3d5024..0000000 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ /dev/null @@ -1,127 +0,0 @@ -from unittest.mock import call, patch, AsyncMock, ANY -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_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'] - }, - 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': workflow_mock.execute_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.write_opc_data, - { - 'opc_output_config': input_data['opc_output_config'], - 'data': workflow_mock.execute_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) - - assert workflow_mock.execute_activity_method.call_count == 2 - 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(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_local_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'] - }, - 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': workflow_mock.execute_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.write_opc_data, - { - 'opc_output_config': input_data['opc_output_config'], - 'data': workflow_mock.execute_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) - - assert workflow_mock.execute_activity_method.call_count == 2 - 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 deleted file mode 100644 index 4318379..0000000 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ /dev/null @@ -1,515 +0,0 @@ -from unittest.mock import AsyncMock, patch, call, ANY -from pytest import fixture, mark -from laborious.activities.activities import Activities -from laborious.workflows.sub_workflows.prediction_process import PredictionProcess - - -@fixture -def prediction_process(): - return PredictionProcess() - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_run(workflow_mock, prediction_process): - prediction_process.path_flag_handler = AsyncMock(return_value=False) - # Arrange - input_data = { - 'data': {'test': 'data'}, - 'schema': 'test_schema', - 'table_name': 'test_table', - 'model_id': 1, - 'input_filters': {'test': 'filter'}, - 'mlflow_transform_filters': {'test': 'filter'}, - 'mlflow_predict_filters': {'test': 'filter'}, - 'model_name': 'test_model_name', - 'model_retention': '30', - 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'}, - } - - # Mock the activity responses - workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp - ('continue', 0.95, "Input data with bad quality"), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data - # mlflow_response_gate (transform) - ('continue', 0.95, "Error"), - # mlflow_content_gate (transform) - ('continue', 0.95, "Transformed data not passed the content filter"), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict - # mlflow_response_gate (predict) - ('continue', 0.95, "Error"), - ] - - # Act - await prediction_process.run(input_data) - - # Assert - assert workflow_mock.execute_local_activity_method.call_count == 7 - - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, {'data': input_data['data']}, - retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.input_gate, { - 'filters': input_data['input_filters'], - 'data': input_data['data'], - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.request_transform, { - 'data': input_data['data'], - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_response_gate, { - 'filters': input_data['mlflow_transform_filters'], - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_content_gate, { - 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.request_predict, { - 'data': 'transformed_data', - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_response_gate, { - 'filters': input_data['mlflow_predict_filters'], - 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, - 'type': 'predict', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - - workflow_mock.execute_child_workflow.assert_called_once_with( - 'format_and_export_prediction', - { - 'path_flag': 'continue', - 'data': 'predicted_data', - 'prediction_confidence': 0.95, - 'timestamp': '2024-01-01', - 'model_id': 1, - 'model_name': 'test_model_name', - 'model_retention': '30', - 'opc_output_config': input_data['opc_output_config'], - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'comment': 'Error' - } - ) - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_run_stop_at_input_gate(workflow_mock, prediction_process): - prediction_process.path_flag_handler = AsyncMock(return_value=True) - # Arrange - input_data = { - 'data': {'test': 'data'}, - 'schema': 'test_schema', - 'table_name': 'test_table', - 'model_id': 1, - 'input_filters': {'test': 'filter'}, - 'mlflow_transform_filters': {'test': 'filter'}, - 'mlflow_predict_filters': {'test': 'filter'}, - 'model_name': 'test_model_name', - 'model_retention': '30', - 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'} - } - - # Mock the activity responses - workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp - ('stop', 0.95, "Input data with bad quality"), # input_gate - ] - - # Act - await prediction_process.run(input_data) - - # Assert - assert workflow_mock.execute_local_activity_method.call_count == 2 - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, { - 'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY), - call(Activities.input_gate, { - 'filters': input_data['input_filters'], - 'data': input_data['data'], - 'path_priority': input_data['path_priority']}, retry_policy=ANY, start_to_close_timeout=ANY) - ]) - workflow_mock.execute_child_workflow.assert_not_called() - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process): - prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True]) - # Arrange - input_data = { - 'data': {'test': 'data'}, - 'schema': 'test_schema', - 'table_name': 'test_table', - 'model_id': 1, - 'input_filters': {'test': 'filter'}, - 'mlflow_transform_filters': {'test': 'filter'}, - 'mlflow_predict_filters': {'test': 'filter'}, - 'model_name': 'test_model_name', - 'model_retention': '30', - 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'} - } - - # Mock the activity responses - workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp - ('repeat', 0.95, "Input data with bad quality"), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data - ('continue', 0.95, "Error"), # mlflow_response_gate (transform) - ] - - # Act - await prediction_process.run(input_data) - - # Assert - assert workflow_mock.execute_local_activity_method.call_count == 4 - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, {'data': input_data['data']}, - retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.input_gate, { - 'filters': input_data['input_filters'], - 'data': input_data['data'], - 'path_priority': input_data['path_priority']}, - retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.request_transform, { - 'data': input_data['data'], - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention']}, - retry_policy=ANY, start_to_close_timeout=ANY) - ]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_response_gate, { - 'filters': input_data['mlflow_transform_filters'], - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY) - ]) - workflow_mock.execute_child_workflow.assert_not_called() - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process): - prediction_process.path_flag_handler = AsyncMock( - side_effect=[False, False, True]) - # Arrange - input_data = { - 'data': {'test': 'data'}, - 'schema': 'test_schema', - 'table_name': 'test_table', - 'model_id': 1, - 'input_filters': {'test': 'filter'}, - 'mlflow_transform_filters': {'test': 'filter'}, - 'mlflow_predict_filters': {'test': 'filter'}, - 'model_name': 'test_model_name', - 'model_retention': '30', - 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'} - } - - # Mock the activity responses - workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp - ('continue', 0.95, "Input data with bad quality"), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data - # mlflow_response_gate (transform) - ('continue', 0.95, "Error"), - # mlflow_content_gate (transform) - ('continue', 0.95, "Transformed data not passed the content filter"), - ] - - # Act - await prediction_process.run(input_data) - - # Assert - assert workflow_mock.execute_local_activity_method.call_count == 5 - - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, {'data': input_data['data']}, - retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.input_gate, { - 'filters': input_data['input_filters'], - 'data': input_data['data'], - 'path_priority': input_data['path_priority']}, - retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.request_transform, { - 'data': input_data['data'], - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_response_gate, { - 'filters': input_data['mlflow_transform_filters'], - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_content_gate, { - 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_child_workflow.assert_not_called() - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process): - prediction_process.path_flag_handler = AsyncMock( - side_effect=[False, False, False, True]) - # Arrange - input_data = { - 'data': {'test': 'data'}, - 'schema': 'test_schema', - 'table_name': 'test_table', - 'model_id': 1, - 'input_filters': {'test': 'filter'}, - 'mlflow_transform_filters': {'test': 'filter'}, - 'mlflow_predict_filters': {'test': 'filter'}, - 'model_name': 'test_model_name', - 'model_retention': '30', - 'path_priority': ['continue', 'repeat', 'stop'], - 'opc_output_config': {'test': 'config'} - } - - # Mock the activity responses - workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp - ('continue', 0.95, "Input data with bad quality"), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data - # mlflow_response_gate (transform) - ('continue', 0.95, "Error"), - # mlflow_content_gate (transform) - ('continue', 0.95, "Transformed data not passed the content filter"), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict - ('continue', 0.95, "Error"), # mlflow_response_gate (predict) - ] - - # Act - await prediction_process.run(input_data) - - # Assert - assert workflow_mock.execute_local_activity_method.call_count == 7 - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, {'data': input_data['data']}, - retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.input_gate, { - 'filters': input_data['input_filters'], - 'data': input_data['data'], - 'path_priority': input_data['path_priority']}, - retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.request_transform, { - 'data': input_data['data'], - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_response_gate, { - 'filters': input_data['mlflow_transform_filters'], - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_content_gate, { - 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', - 'type': 'transform', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.request_predict, { - 'data': 'transformed_data', - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.mlflow_response_gate, { - 'filters': input_data['mlflow_predict_filters'], - 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, - 'type': 'predict', - 'path_priority': input_data['path_priority'] - }, retry_policy=ANY, start_to_close_timeout=ANY)]) - workflow_mock.execute_child_workflow.assert_not_called() - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_path_flag_handler_stop(workflow_mock, prediction_process): - # Arrange - data = {'test': 'data'} - path_flag = 'STOP' - confidence = 0.95 - schema = 'test_schema' - table_name = 'test_table' - model = 'test_model' - last_timestamp = '2024-01-01' - model_name = 'test_model_name' - model_retention = '30' - - # Act - result = await prediction_process.path_flag_handler( - data, path_flag, { - 'schema': schema, - 'table_name': table_name, - 'model_id': model, - 'last_timestamp': last_timestamp, - 'model_name': model_name, - 'model_retention': model_retention - }, confidence, last_timestamp, "" - ) - - # Assert - assert result is True - workflow_mock.execute_local_activity_method.assert_not_called() - workflow_mock.execute_child_workflow.assert_not_called() - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_path_flag_handler_repeat(workflow_mock, prediction_process): - # Arrange - data = {'test': 'data'} - path_flag = 'repeat' - confidence = 0.95 - schema = 'test_schema' - table_name = 'test_table' - model = 'test_model' - last_timestamp = '2024-01-01' - model_name = 'test_model_name' - model_retention = '30' - - # Act - result = await prediction_process.path_flag_handler( - data, path_flag, { - 'schema': schema, - 'table_name': table_name, - 'model_id': model, - 'last_timestamp': last_timestamp, - 'model_name': model_name, - 'model_retention': model_retention - }, confidence, last_timestamp, "" - ) - - # Assert - assert result is True - workflow_mock.execute_activity_method.assert_called_once_with( - Activities.repeat_last_prediction, - { - 'schema': schema, - 'table_name': table_name, - 'model_id': model - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - workflow_mock.execute_child_workflow.assert_not_called() - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_path_flag_handler_continue(workflow_mock, prediction_process): - # Arrange - data = {'test': 'data'} - path_flag = 'CONTINUE' - confidence = 0.95 - schema = 'test_schema' - table_name = 'test_table' - model = 'test_model' - last_timestamp = '2024-01-01' - model_name = 'test_model_name' - model_retention = '30' - - # Act - result = await prediction_process.path_flag_handler( - data, path_flag, { - 'schema': schema, - 'table_name': table_name, - 'model_id': model, - 'last_timestamp': last_timestamp, - 'model_name': model_name, - 'model_retention': model_retention, - 'opc_output_config': {'test': 'config'} - }, confidence, last_timestamp, 'Prediction Process' - ) - - # Assert - assert result is True - workflow_mock.execute_activity_method.assert_not_called() - workflow_mock.execute_child_workflow.assert_called_once_with( - '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, - 'schema': schema, - 'table_name': table_name, - 'comment': 'Prediction Process', - 'opc_output_config': {'test': 'config'} - } - ) - - -@mark.asyncio -@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) -async def test_path_flag_handler_unknown(workflow_mock, prediction_process): - # Arrange - data = {'test': 'data'} - path_flag = 'unknown' - confidence = 0.95 - schema = 'test_schema' - table_name = 'test_table' - model = 'test_model' - last_timestamp = '2024-01-01' - model_name = 'test_model_name' - model_retention = '30' - - # Act - result = await prediction_process.path_flag_handler( - data, path_flag, { - 'schema': schema, - 'table_name': table_name, - 'model_id': model, - 'last_timestamp': last_timestamp, - 'model_name': model_name, - 'model_retention': model_retention, - 'opc_output_config': {'test': 'config'} - }, confidence, last_timestamp, "" - ) - - # Assert - assert result is False - workflow_mock.execute_activity_method.assert_not_called() - workflow_mock.execute_child_workflow.assert_not_called() diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py deleted file mode 100644 index 0ca45e1..0000000 --- a/tests/laborious/workflows/test_predictions_batch.py +++ /dev/null @@ -1,81 +0,0 @@ -from unittest.mock import AsyncMock, call, patch, ANY -from pytest import fixture, mark -from laborious.activities.activities import Activities -from laborious.workflows.predictions_batch import PredictionsBatch - - -@fixture -def predictions_batch() -> PredictionsBatch: - return PredictionsBatch() - - -@mark.asyncio -@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock) -async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch): - workflow_mock.execute_local_activity_method.return_value = { - 'data': 'test_data' - } - input_data = { - 'schedule_name': 'test_schedule', - 'model_name': 'test_model', - 'model_id': 'test_model_id', - 'query': 'SELECT * FROM test', - 'schema': 'test_schema', - 'table_name': 'test_table', - 'opc_output_config': 'test_opc_output_config' - } - - await predictions_batch.run(input_data) - - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.prepare_activity, - { - 'schedule_name': input_data['schedule_name'], - 'model_name': input_data['model_name'], - 'model_id': input_data['model_id'], - 'workflow_name': 'predictions_batch' - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) - - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.load_custom_query, - input_data['query'], - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) - prediction_input = { - 'data': {'data': 'test_data'}, - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'model_id': input_data['model_id'], - 'model_name': input_data['model_name'], - 'input_filters': input_data.get('input_filters', { - 'EMPTY_DATA': { - 'POLICY': 'STOP' - } - }), - 'mlflow_transform_filters': input_data.get('mlflow_transform_filters', { - 'API_ERROR': { - 'POLICY': 'STOP' - } - }), - 'mlflow_predict_filters': input_data.get('mlflow_predict_filters', { - 'API_ERROR': { - 'POLICY': 'STOP' - } - }), - 'model_retention': input_data.get('model_retention', 60), - 'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), - 'opc_output_config': input_data.get('opc_output_config', {}) - } - - workflow_mock.execute_child_workflow.assert_has_calls([ - call( - 'prediction_process', prediction_input) - ]) diff --git a/tests/laborious/__init__.py b/tests/orchestrator/__init__.py similarity index 100% rename from tests/laborious/__init__.py rename to tests/orchestrator/__init__.py diff --git a/tests/orchestrator/activities/test_activities.py b/tests/orchestrator/activities/test_activities.py new file mode 100644 index 0000000..14d857c --- /dev/null +++ b/tests/orchestrator/activities/test_activities.py @@ -0,0 +1,116 @@ +from unittest.mock import patch, MagicMock, ANY +from pytest import mark +from orchestrator.activities.activities import Activities +from orchestrator.activities.couchbase import Couchbase +from orchestrator.activities.temporal_manager import TemporalManager +from orchestrator.activities.slot_manager import SlotManager + + +@patch('orchestrator.activities.couchbase.Couchbase.__init__') +@patch('orchestrator.activities.temporal_manager.TemporalManager.__init__') +@patch('orchestrator.activities.slot_manager.SlotManager.__init__') +def test___init__(mock_slot_manager_init, mock_temporal_manager_init, + mock_couchbase_init): + + couchbase_config = { + 'connection_string': 'couchbase://localhost', + 'username': 'admin', + 'password': 'password' + } + + redis_config = { + 'host': 'localhost', + 'port': 6379, + 'username': 'admin', + 'password': 'password' + } + + temporal_client = MagicMock() + logger = MagicMock() + notification_handler = MagicMock() + + activities = Activities( + temporal_client=temporal_client, + couchbase_config=couchbase_config, + redis_config=redis_config, + logger=logger, + notification_handler=notification_handler + ) + + assert isinstance(activities, Activities) + assert isinstance(activities, Couchbase) + assert isinstance(activities, TemporalManager) + assert isinstance(activities, SlotManager) + + mock_slot_manager_init.assert_called_once_with( + ANY, + host="localhost", + port=6379, + username="admin", + password="password", + logger=logger, + notification_handler=notification_handler + ) + + mock_couchbase_init.assert_called_once_with( + ANY, + connection_string=couchbase_config['connection_string'], + username=couchbase_config['username'], + password=couchbase_config['password'], + logger=logger, + notification_handler=notification_handler + ) + + mock_temporal_manager_init.assert_called_once_with( + ANY, + temporal_client=temporal_client, + logger=logger, + notification_handler=notification_handler + ) + + +@mark.asyncio +@patch('orchestrator.activities.couchbase.Cluster') +async def test_prepare_activity(_mock_cluster): + couchbase_config = { + 'connection_string': 'couchbase://localhost', + 'username': 'admin', + 'password': 'password' + } + + redis_config = { + 'host': 'localhost', + 'port': 6379, + 'username': 'admin', + 'password': 'password' + } + + temporal_client = MagicMock() + logger = MagicMock() + notification_handler = MagicMock() + + activities = Activities( + temporal_client=temporal_client, + couchbase_config=couchbase_config, + redis_config=redis_config, + logger=logger, + notification_handler=notification_handler + ) + + input_data = { + 'workflow_name': 'test-workflow-name', + 'schedule_name': 'test-schedule-name', + 'model_name': 'test-model-name', + 'model_id': 'test-model-id' + } + + await activities.prepare_activity(input_data) + + assert activities.notification_handler.base_notification.pipeline == input_data[ + 'workflow_name'] + assert activities.notification_handler.base_notification.trigger == input_data[ + 'schedule_name'] + assert activities.notification_handler.base_notification.model_name == input_data[ + 'model_name'] + assert activities.notification_handler.base_notification.model_id == input_data[ + 'model_id'] diff --git a/tests/orchestrator/activities/test_couchbase.py b/tests/orchestrator/activities/test_couchbase.py new file mode 100644 index 0000000..84f4da7 --- /dev/null +++ b/tests/orchestrator/activities/test_couchbase.py @@ -0,0 +1,56 @@ +from unittest.mock import MagicMock, patch, ANY +from pytest import fixture, mark, raises +from sientia_do.notifications.models import NotificationLevel +from orchestrator.activities.couchbase import Couchbase + + +@fixture +@patch("orchestrator.activities.couchbase.Cluster") +def couchbase(_cluster_mock): + return Couchbase( + connection_string="couchbase://localhost", + username="admin", + password="password", + logger=MagicMock(), + notification_handler=MagicMock(), + ) + + +@mark.asyncio +async def test_load_query_from_couchbase_success(couchbase): + couchbase.cluster.query.return_value.rows.return_value = [ + {"id": "1", "name": "test"}, + {"id": "2", "name": "test2"}, + ] + query = "SELECT * FROM bucket" + + result = await couchbase.load_query_from_couchbase({ + "query": query, + }) + + assert result == [ + {"id": "1", "name": "test"}, + {"id": "2", "name": "test2"}, + ] + couchbase.cluster.query.assert_called_once_with(query) + couchbase.notification_handler.build_and_send_notification.assert_not_called() + + +@mark.asyncio +async def test_load_query_from_couchbase_failure(couchbase): + couchbase.cluster.query.side_effect = Exception("Test error") + query = "SELECT * FROM bucket" + + with raises(Exception): + await couchbase.load_query_from_couchbase({ + "query": query, + }) + + couchbase.cluster.query.assert_called_once_with(query) + couchbase.notification_handler.build_and_send_notification.assert_called_once_with( + notification_id="COUCHBASE_LOAD_QUERY_ERROR", + message="Failed to execute couchbase query: Test error", + block="load_query_from_couchbase", + level=NotificationLevel.ERROR, + attachment_content=ANY, + ) diff --git a/tests/orchestrator/activities/test_slot_manager.py b/tests/orchestrator/activities/test_slot_manager.py new file mode 100644 index 0000000..cca97ca --- /dev/null +++ b/tests/orchestrator/activities/test_slot_manager.py @@ -0,0 +1,57 @@ +from unittest.mock import MagicMock, patch +from pytest import mark, fixture +from orchestrator.activities.slot_manager import SlotManager + + +@fixture +@patch("orchestrator.activities.slot_manager.Redis.__init__") +def slot_manager(_redis_mock): + + slot_manager = SlotManager( + host="localhost", + port=6379, + username="admin", + password="password", + logger=MagicMock(), + notification_handler=MagicMock() + ) + + slot_manager.redis_client = MagicMock() + slot_manager.logger = MagicMock() + slot_manager.notification_handler = MagicMock() + + return slot_manager + + +@mark.asyncio +async def test_load_opc_slots_no_slot_keys(slot_manager): + slot_manager.redis_client.keys.return_value = [] + assert await slot_manager.load_opc_slots() == {} + + +@mark.asyncio +async def test_load_opc_slots(slot_manager): + slot_manager.redis_client.keys.return_value = [ + b"slot:opc_tags:1", b"slot:opc_tags:2", b"slot:opc_tags:3"] + + slot_manager.redis_client.mget.return_value = [ + b"value1", "value2", None] + + response = await slot_manager.load_opc_slots() + + assert response == { + "slot:opc_tags:1": "value1", + "slot:opc_tags:2": "value2", + "slot:opc_tags:3": None + } + + +@mark.asyncio +async def test_load_active_ingestors(slot_manager): + slot_manager.redis_client.keys.return_value = [ + b"heartbeat:ingestor:1", b"heartbeat:ingestor:2", b"heartbeat:ingestor:3"] + + response = await slot_manager.load_active_ingestors() + + assert response == ["heartbeat:ingestor:1", + "heartbeat:ingestor:2", "heartbeat:ingestor:3"] diff --git a/tests/orchestrator/activities/test_temporal_manager.py b/tests/orchestrator/activities/test_temporal_manager.py new file mode 100644 index 0000000..6808bdc --- /dev/null +++ b/tests/orchestrator/activities/test_temporal_manager.py @@ -0,0 +1,80 @@ +from unittest.mock import MagicMock, patch, AsyncMock +import base64 +import json +from pytest import fixture, mark +from orchestrator.activities.temporal_manager import TemporalManager + + +@fixture +def temporal_manager(): + return TemporalManager( + temporal_client=MagicMock(), + logger=MagicMock(), + notification_handler=MagicMock() + ) + + +@mark.asyncio +@patch("orchestrator.activities.temporal_manager.MessageToDict", + return_value={"data": base64.b64encode(json.dumps({"test": "test"}).encode('utf-8'))}) +async def test_load_schedule(_mock_message_to_dict, temporal_manager): + # Create async iterator mock + async def async_iter(): + yield MagicMock( + id="test-schedule-id", + search_attributes={ + "Orchestrated": ["true"] + } + ) + yield MagicMock( + id="test-schedule-id-2", + search_attributes={ + "Attr": ["false"] + } + ) + yield MagicMock( + id="test-schedule-id-3", + search_attributes={ + "Attr": ["false"] + } + ) + + temporal_manager.temporal_client.list_schedules = AsyncMock( + return_value=async_iter()) + temporal_manager.temporal_client.get_schedule.return_value = MagicMock( + describe=AsyncMock( + return_value=MagicMock( + schedule=MagicMock( + action=MagicMock( + args=[ + MagicMock( + data=base64.b64encode(json.dumps( + {"test": "test"}).encode('utf-8')) + ) + ] + ) + ) + ) + ) + ) + temporal_manager.temporal_client.get_schedule.return_value.describe \ + .return_value.schedule.spec = MagicMock( + intervals=[ + MagicMock( + every=MagicMock( + seconds=60 + ) + ) + ] + ) + + response = await temporal_manager.load_schedule() + + temporal_manager.temporal_client.list_schedules.assert_called_once() + assert response == { + "test-schedule-id": { + "frequency": 60, + "data": {"test": "test"}, + "handle": temporal_manager.temporal_client.get_schedule.return_value + } + } diff --git a/tests/orchestrator/workflows/test_orchestrator.py b/tests/orchestrator/workflows/test_orchestrator.py new file mode 100644 index 0000000..a9d80f8 --- /dev/null +++ b/tests/orchestrator/workflows/test_orchestrator.py @@ -0,0 +1,81 @@ +from unittest.mock import AsyncMock, patch, ANY, call +from pytest import fixture, mark +from orchestrator.workflows.orchestrator import Orchestrator +from orchestrator.activities.activities import Activities + + +@fixture +def orchestrator(): + return Orchestrator() + + +@mark.asyncio +@patch("orchestrator.workflows.orchestrator.workflow", new_callable=AsyncMock) +async def test_run(workflow_mock, orchestrator): + input_data = { + "pipelines_query": "SELECT * FROM bucket", + "opc_servers_query": "SELECT * FROM servers", + "schedule_name": "test-schedule-name", + } + + await orchestrator.run(input_data) + + workflow_mock.execute_local_activity_method.assert_has_calls([ + call( + Activities.prepare_activity, + { + "workflow_name": "orchestrator", + "schedule_name": "test-schedule-name", + "model_name": "-", + "model_id": "-" + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + + workflow_mock.execute_local_activity_method.assert_has_calls([ + call( + Activities.load_query_from_couchbase, + { + "query": input_data["pipelines_query"] + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + + workflow_mock.execute_local_activity_method.assert_has_calls([ + call( + Activities.load_query_from_couchbase, + { + "query": input_data["opc_servers_query"] + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + + workflow_mock.execute_local_activity_method.assert_has_calls([ + call( + Activities.load_schedule, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + + workflow_mock.execute_local_activity_method.assert_has_calls([ + call( + Activities.load_opc_slots, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + + workflow_mock.execute_local_activity_method.assert_has_calls([ + call( + Activities.load_active_ingestors, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ])