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

SIENTIAPDE-994

Browse Source 
Refactor and enhance the laborious workflow and utilities

- Removed outdated test file `test_predictions_batch.py` from workflows.
- Added `input_sample.json` for standardized input configuration.
- Introduced `connectors_config.py` to manage database and service configurations.
- Implemented a logging utility in `logger.py` for consistent logging across the application.
- Created `policies.py` to define retry policies for workflows.
- Developed comprehensive tests for `MLFlowRepository` in `test_model_repository.py`.
- Added extensive tests for `OpcRepository` in `test_opc_repository.py`.
- Updated `test_predictions_batch.py` to reflect new workflow structure and testing methodology.
This commit is contained in:
vitor-aignosi
2025-05-23 17:34:47 -03:00
parent 5fe552410b
commit 67fe4afaa6
30 changed files with 1385 additions and 765 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
input_sample.json Normal file
View File

@@ -0,0 +1,34 @@
{
"schedule_name": "scouter-opcua-pipeline",
"model_name": "Demo Model",
"model_id": 1,
"query": "SELECT * FROM sientia_data.laborious_data order by \"timestamp\" desc limit 30;",
"schema": "sientia_data",
"table_name": "predictions",
"retention_time": 3600,
"model_retention": 120,
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
"input_filters": {
"SPECIFIC_VARIABLES_NULL_VALUES": {
"POLICY": "STOP",
"VARIABLES": ["Counter"]
},
"EMPTY_DATA": {
"POLICY": "STOP"
}
},
"mlflow_transform_filters": {
"API_ERROR": {
"POLICY": "CONTINUE"
},
"NAN_VALUES": {
"POLICY": "CONTINUE"
}
},
"mlflow_predict_filters": {
"API_ERROR": {
"POLICY": "CONTINUE"
}
},
"opc_output_config": {}
}

11
laborious/activities/activities.py
View File

@@ -40,15 +40,10 @@ class Activities(Postgres, MLFlow, Gates, OPC):
notification_handler=notification_handler) notification_handler=notification_handler)
OPC.__init__(self, OPC.__init__(self,
name=opc_config['name'], opc_servers=opc_config,
url=opc_config['url'],
server_uri=opc_config['server_uri'],
cert_path=opc_config['cert_path'],
private_key_path=opc_config['private_key_path'],
server_cert_path=opc_config['server_cert_path'],
logger=logger, logger=logger,
notification_handler=notification_handler) notification_handler=notification_handler)
@activity.defn(name="prepare_activity") @activity.defn(name="prepare_activity")
def prepare_activity(self, input_data: dict[str, Any]): async def prepare_activity(self, input_data: dict[str, Any]):
super().prepare_activity(input_data) await super().prepare_activity(input_data)

3
laborious/activities/base.py
View File

@@ -1,3 +1,4 @@
from typing import Any
from logging import Logger from logging import Logger
from temporalio import activity from temporalio import activity
from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.handlers import NotificationHandler
@@ -9,7 +10,7 @@ class BaseActivity:
self.notification_handler = notification_handler self.notification_handler = notification_handler
@activity.defn(name="prepare_activity") @activity.defn(name="prepare_activity")
def prepare_activity(self, input_data: dict[str, Any]): async def prepare_activity(self, input_data: dict[str, Any]):
""" """
Prepare the activity for the notification handler. Prepare the activity for the notification handler.

115
laborious/activities/gates.py
View File

@@ -5,67 +5,85 @@ with workflow.unsafe.imports_passed_through():
import traceback import traceback
from logging import Logger from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.handlers import NotificationHandler
from laborious.activities.base import BaseActivity
from typing import Any
from laborious.utils.filters.conditional_filters import filter_empty_data, filter_specific_variables_null_values
from pandas import DataFrame
from sientia_do.notifications.models import NotificationLevel 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 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 = { input_filter_functions = {
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values, 'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
'EMPTY_DATA': filter_empty_data, 'EMPTY_DATA': filter_empty_data,
'path_confidence': { 'path_confidence': {
'stop': -1, 'STOP': -1,
'continue': 2, 'CONTINUE': 2,
'repeat': -1 'REPEAT': -1
} }
} }
mlflow_response_filter_functions = { mlflow_response_filter_functions = {
'API_ERROR': api_error_filter, 'API_ERROR': api_error_filter,
'path_confidence': { 'path_confidence': {
'stop': -1, 'STOP': -1,
'continue': 10, 'CONTINUE': 10,
'repeat': -1 'REPEAT': -1
}, },
} }
mlflow_content_filter_functions = { mlflow_content_filter_functions = {
'NAN_VALUES': nan_values_filter, 'NAN_VALUES': nan_values_filter,
'path_confidence': { 'path_confidence': {
'stop': -1, 'STOP': -1,
'continue': 18, 'CONTINUE': 18,
'repeat': -1 'REPEAT': -1
} }
} }
class Gates(BaseActivity): class Gates(BaseActivity):
def __init__(self, logger: Logger, notification_handler: NotificationHandler): def __init__(self, logger: Logger, notification_handler: NotificationHandler):
super().__init__(logger, notification_handler) BaseActivity.__init__(self, logger, notification_handler)
@activity.defn(name="input_gate") @activity.defn(name="input_gate")
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str, int]: 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 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. being the policy and the second element being the confidence status.
Args: Args:
input_data (dict): The input data. Contains: input_data (dict): The input data. Contains:
filters (dict): The filters to apply. 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. data (dict[str, Any]): The data to filter.
path_priority (list[str]): The path priority. path_priority (list[str]): The path priority.
Returns: Returns:
tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions. 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'] filters = input_data['filters']
data = DataFrame(input_data['data']) data = DataFrame(input_data['data'])
path_priority = input_data['path_priority'] path_priority = input_data['path_priority']
filter_output = [] filter_output = []
self.logger.debug(f"Input data:\n {data.to_string()}")
self.logger.debug(f"Filters: {filters}")
for fil, config in filters.items(): for fil, config in filters.items():
if fil not in input_filter_functions:
self.logger.error(f"Filter {fil} not found")
continue
try: try:
if input_filter_functions[fil](data, config): if input_filter_functions[fil](data, config):
self.logger.debug(
f"Data not passed the input filter {fil}:{config}")
filter_output.append(config['POLICY']) filter_output.append(config['POLICY'])
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
@@ -79,14 +97,18 @@ class Gates(BaseActivity):
for path_flag in path_priority: for path_flag in path_priority:
if path_flag in filter_output: if path_flag in filter_output:
return path_flag, input_filter_functions['path_confidence'][path_flag] self.logger.debug(f"Input gate result: {path_flag}")
return path_flag, input_filter_functions['path_confidence'][path_flag], \
"Input data with bad quality"
return None, 0 self.logger.debug("Nothing was filtered by the input gate")
return None, 0, ""
@activity.defn(name="mlflow_response_gate") @activity.defn(name="mlflow_response_gate")
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str, int]: 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 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. being the policy and the second element being the confidence status.
Args: Args:
input_data (dict): The input data. Contains: input_data (dict): The input data. Contains:
@@ -95,17 +117,29 @@ class Gates(BaseActivity):
path_priority (list[str]): The path priority list. path_priority (list[str]): The path priority list.
type (str): The type of the gate. type (str): The type of the gate.
Returns: Returns:
tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions. 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'] filters = input_data['filters']
data = input_data['data'] data = input_data['data']
gate_type = input_data['type'] gate_type = input_data['type']
path_priority = input_data['path_priority'] path_priority = input_data['path_priority']
filter_output = [] filter_output = []
self.logger.debug(f"Input data:\n {data}")
self.logger.debug(f"Filters: {filters}")
comments = []
for fil, config in filters.items(): for fil, config in filters.items():
if fil not in mlflow_response_filter_functions:
continue
if mlflow_response_filter_functions[fil](data, config): if mlflow_response_filter_functions[fil](data, config):
filter_output.append(config['POLICY']) filter_output.append(config['POLICY'])
comments.append(data['content']['message'])
self.notification_handler.build_and_send_notification( self.notification_handler.build_and_send_notification(
notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}", notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
message=data['content']['message'], message=data['content']['message'],
@@ -116,14 +150,18 @@ class Gates(BaseActivity):
for path_flag in path_priority: for path_flag in path_priority:
if path_flag in filter_output: if path_flag in filter_output:
return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag] self.logger.debug(f"Mlflow response gate result: {path_flag}")
return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \
", ".join(comments)
return None, 0 self.logger.debug("Nothing was filtered by the mlflow response gate")
return None, 0, ""
@activity.defn(name="mlflow_content_gate") @activity.defn(name="mlflow_content_gate")
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str, int]: 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 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. being the policy and the second element being the confidence status.
Args: Args:
input_data (dict): The input data. Contains: input_data (dict): The input data. Contains:
@@ -132,9 +170,12 @@ class Gates(BaseActivity):
path_priority (list[str]): The path priority list. path_priority (list[str]): The path priority list.
type (str): The type of the gate. type (str): The type of the gate.
Returns: Returns:
tuple[str, int]: (policy, confidence) based in priority list and filter configuration and functions. 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'] filters = input_data['filters']
data = DataFrame(input_data['data']) data = DataFrame(input_data['data'])
gate_type = input_data['type'] gate_type = input_data['type']
@@ -142,7 +183,12 @@ class Gates(BaseActivity):
filter_output = [] filter_output = []
self.logger.debug(f"Input data:\n {data}")
self.logger.debug(f"Filters: {filters}")
for fil, config in filters.items(): for fil, config in filters.items():
if fil not in mlflow_content_filter_functions:
continue
if mlflow_content_filter_functions[fil](data, config): if mlflow_content_filter_functions[fil](data, config):
filter_output.append(config['POLICY']) filter_output.append(config['POLICY'])
self.notification_handler.build_and_send_notification( self.notification_handler.build_and_send_notification(
@@ -155,9 +201,12 @@ class Gates(BaseActivity):
for path_flag in path_priority: for path_flag in path_priority:
if path_flag in filter_output: if path_flag in filter_output:
return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag] 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"
return None, 0 self.logger.debug("Nothing was filtered by the mlflow content gate")
return None, 0, ""
@activity.defn(name="format_prediction") @activity.defn(name="format_prediction")
async def format_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]: async def format_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]:
@@ -172,12 +221,14 @@ class Gates(BaseActivity):
Returns: Returns:
dict: The formatted data. dict: The formatted data.
""" """
self.logger.debug("Formatting prediction...")
data = DataFrame(input_data['data']) data = DataFrame(input_data['data'])
data['timestamp'] = input_data['timestamp'] data['timestamp'] = input_data['timestamp']
data['model_id'] = input_data['model_id'] data['model_id'] = input_data['model_id']
data['prediction_confidence'] = input_data['prediction_confidence'] data['prediction_confidence'] = input_data['prediction_confidence']
data['prediction_status'] = 'Good' data['prediction_status'] = 'Good'
data['comment'] = "" data['comments'] = ""
data.sort_values(by='timestamp', inplace=True) data.sort_values(by='timestamp', inplace=True)
return data.to_dict() return data.to_dict()
@@ -198,6 +249,8 @@ class Gates(BaseActivity):
dict: The formatted data. dict: The formatted data.
""" """
self.logger.debug("Formatting default prediction...")
return DataFrame({ return DataFrame({
'prediction': [0], 'prediction': [0],
'response_time': [0], 'response_time': [0],
@@ -205,7 +258,7 @@ class Gates(BaseActivity):
'model_id': [input_data['model_id']], 'model_id': [input_data['model_id']],
'prediction_confidence': [input_data['prediction_confidence']], 'prediction_confidence': [input_data['prediction_confidence']],
'prediction_status': ['Bad'], 'prediction_status': ['Bad'],
'comment': [input_data['comment']] 'comments': [input_data['comment']]
}).to_dict() }).to_dict()
@activity.defn(name="get_last_timestamp") @activity.defn(name="get_last_timestamp")
@@ -219,4 +272,6 @@ class Gates(BaseActivity):
str: The last timestamp of the data. str: The last timestamp of the data.
""" """
data = DataFrame(input_data['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()) return max(data['timestamp'].values.tolist())

11
laborious/activities/mlflow.py
View File

@@ -5,17 +5,16 @@ from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from laborious.activities.base import BaseActivity from laborious.activities.base import BaseActivity
from laborious.utils.repository.model_repository import MLFlowRepository
from typing import Any from typing import Any
from logging import Logger from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.handlers import NotificationHandler
from laborious.utils.repository.model_repository import MLFlowRepository
from sientia_do.notifications.models import NotificationLevel
class MLFlow(BaseActivity): class MLFlow(BaseActivity):
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str, def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
mlflow_password: str, logger: Logger, notification_handler: NotificationHandler): mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
super().__init__(logger, notification_handler) BaseActivity.__init__(self, logger, notification_handler)
self.mlflow_host = mlflow_host self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username self.mlflow_username = mlflow_username
@@ -33,7 +32,7 @@ class MLFlow(BaseActivity):
input_data (dict): The input data. Contains: input_data (dict): The input data. Contains:
data (dict[str, Any]): The data to transform. data (dict[str, Any]): The data to transform.
model_name (str): The name of the model. model_name (str): The name of the model.
model_retention (int): The retention of the model. model_retention (int): The retention of the model in minutes.
Returns: Returns:
dict[str, Any]: The transformed data. dict[str, Any]: The transformed data.
""" """
@@ -54,6 +53,8 @@ class MLFlow(BaseActivity):
response_data = self.model_monitoring_repository.transform( response_data = self.model_monitoring_repository.transform(
model_name, data, model_retention) model_name, data, model_retention)
self.logger.debug(response_data)
return response_data return response_data
@activity.defn(name="request_predict") @activity.defn(name="request_predict")
@@ -80,4 +81,6 @@ class MLFlow(BaseActivity):
response_data = self.model_monitoring_repository.predict( response_data = self.model_monitoring_repository.predict(
model_name, data, model_retention) model_name, data, model_retention)
self.logger.debug(response_data)
return response_data return response_data

115
laborious/activities/opc.py
View File

@@ -4,40 +4,55 @@ from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from logging import Logger from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.base import BaseActivity from laborious.activities.base import BaseActivity
from laborious.utils.repository.opc_repository import OpcRepository from laborious.utils.repository.opc_repository import OpcRepository
from typing import Any from typing import Any
from sientia_do.notifications.models import NotificationLevel
import traceback import traceback
from pandas import DataFrame from pandas import DataFrame
class OPC(BaseActivity): class OPC(BaseActivity):
def __init__(self, def __init__(self, opc_servers: dict[str, dict[str, Any]],
name: str, url: str, server_uri: str,
cert_path: str, private_key_path: str, server_cert_path: str,
logger: Logger, notification_handler: NotificationHandler): logger: Logger, notification_handler: NotificationHandler):
self.logger = logger self.logger = logger
self.notification_handler = notification_handler self.notification_handler = notification_handler
self.name = name self.opc_servers = opc_servers
self.url = url
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.opc_repository = OpcRepository( self.opc_repository = {}
name=self.name, for name, server in opc_servers.items():
url=self.url, self.opc_repository[name] = OpcRepository(
logger=self.logger, name=name,
server_uri=self.server_uri, url=server['url'],
cert_path=self.cert_path, logger=self.logger,
private_key_path=self.private_key_path, server_uri=server['server_uri'],
server_cert_path=self.server_cert_path 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()
self.opc_repository.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') @activity.defn(name='write_opc_data')
async def write_opc_data(self, input_data: dict[str, Any]): async def write_opc_data(self, input_data: dict[str, Any]):
@@ -47,46 +62,40 @@ class OPC(BaseActivity):
Args: Args:
input_data (dict[str, Any]): The input data. Contains the following keys: 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. - data (dict[str, Any]): The dataframe that contains the data to write
opc_servers (list[str]): The OPC servers to write to. to the OPC servers.
opc_output_config (dict[str, Any]): The OPC writing configuration. Contains: - 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. 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. confidence_tags (dict[str, Any]): The tags to write to the OPC servers.
Returns: Returns:
""" """
self.logger.debug("Writing data to OPC servers...")
data = DataFrame(input_data['data']) data = DataFrame(input_data['data'])
_opc_servers = input_data['opc_servers']
opc_output_config = input_data['opc_output_config'] opc_output_config = input_data['opc_output_config']
self.logger.debug(data)
if 'prediction_tags' in opc_output_config: for server, config in opc_output_config.items():
for tag, config in opc_output_config['prediction_tags'].items(): if self.opc_repository.get(server) is None:
try: self.logger.error(f"OPC server {server} not found")
self.opc_repository.write_data( continue
tag, data.head(1)['prediction'].values[0], config['data_type'])
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id="WRITE_OPC_PREDICTION_ERROR",
message=f"Error writing data to OPC server: {e}",
block="write_opc_data",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
if 'confidence_tags' in opc_output_config: if 'prediction_tags' in config:
for tag, config in opc_output_config['confidence_tags'].items(): for tag, tag_config in config['prediction_tags'].items():
try: self.write_data(
self.opc_repository.write_data( server=server,
tag, data.head(1)['prediction_confidence'].values[0], config['data_type']) tag=tag,
except Exception as e: data=data.head(1)['prediction'].values[0],
trace = traceback.format_exc() data_type=tag_config['data_type'],
self.notification_handler.build_and_send_notification( tag_type='prediction'
notification_id="WRITE_OPC_CONFIDENCE_ERROR", )
message=f"Error writing data to OPC server: {e}", if 'confidence_tags' in config:
block="write_opc_data", for tag, tag_config in config['confidence_tags'].items():
level=NotificationLevel.ERROR, self.write_data(
attachment_content=trace server=server,
tag=tag,
data=data.head(1)['prediction_confidence'].values[0],
data_type=tag_config['data_type'],
tag_type='confidence'
) )
self.logger.error(trace)

141
laborious/activities/postgres.py
View File

@@ -3,6 +3,9 @@ from temporalio import workflow, activity
from laborious.activities.base import BaseActivity from laborious.activities.base import BaseActivity
with workflow.unsafe.imports_passed_through(): 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 psycopg2.pool import ThreadedConnectionPool
from pandas import read_sql_query, DataFrame from pandas import read_sql_query, DataFrame
from logging import Logger from logging import Logger
@@ -22,19 +25,20 @@ class Postgres(BaseActivity):
self.password = password self.password = password
self.dbname = dbname self.dbname = dbname
self.pool = ThreadedConnectionPool( # Create SQLAlchemy engine with connection pooling
minconn=min_connections, self.engine = create_engine(
maxconn=max_connections, f'postgresql://{user}:{password}@{host}:{port}/{dbname}',
host=self.host, poolclass=QueuePool,
port=self.port, pool_size=min_connections,
user=self.user, max_overflow=max_connections - min_connections,
password=self.password, pool_pre_ping=True
dbname=self.dbname) )
self.session_factory = sessionmaker(bind=self.engine)
super().__init__(logger, notification_handler) BaseActivity.__init__(self, logger, notification_handler)
def close(self): def close(self):
self.pool.closeall() self.engine.dispose()
def __del__(self): def __del__(self):
self.close() self.close()
@@ -52,28 +56,36 @@ class Postgres(BaseActivity):
""" """
self.logger.info(f"Fetching data from query: {query}") self.logger.info(f"Fetching data from query: {query}")
conn = self.pool.getconn() data = None
try: with self.session_factory() as session:
data = read_sql_query(query, conn) try:
data = read_sql_query(query, self.engine)
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.notification_handler.build_and_send_notification( self.notification_handler.build_and_send_notification(
notification_id="ERROR_LOADING_CUSTOM_QUERY", notification_id="ERROR_LOADING_CUSTOM_QUERY",
message=f"Error fetching data from query: {e}", message=f"Error fetching data from query: {e}",
block="load_custom_query", block="load_custom_query",
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace
) )
self.logger.error(trace) self.logger.error(trace)
return {}
finally:
session.close()
if data is None:
return {} return {}
finally:
self.pool.putconn(conn) # 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.info(f"Fetched {len(data)} rows")
self.logger.debug(f"Data: {data.to_string()}") self.logger.debug(f"Data: \n{data.to_string()}")
return data.to_dict() return data.to_dict()
@@ -86,7 +98,7 @@ class Postgres(BaseActivity):
query_items (dict[str, str]): The query items. Contains: query_items (dict[str, str]): The query items. Contains:
schema (str): The schema of the table. schema (str): The schema of the table.
table_name (str): The name of the table. table_name (str): The name of the table.
model (str): The model to repeat the prediction for. model (int): The model to repeat the prediction for.
Returns: Returns:
None None
@@ -106,28 +118,25 @@ class Postgres(BaseActivity):
self.logger.info(f"Repeating last prediction for model {model}") self.logger.info(f"Repeating last prediction for model {model}")
self.logger.debug(f"Query: {repeat_query}") self.logger.debug(f"Query: {repeat_query}")
conn = self.pool.getconn() with self.session_factory() as session:
try:
session.execute(repeat_query)
session.commit()
try: except Exception as e:
cursor = conn.cursor() trace = traceback.format_exc()
cursor.execute(repeat_query) self.notification_handler.build_and_send_notification(
conn.commit() notification_id="ERROR_REPEATING_LAST_PREDICTION",
cursor.close() message=f"Error repeating last prediction: {e}",
block="repeat_last_prediction",
level=NotificationLevel.ERROR,
attachment_content=trace
)
except Exception as e: self.logger.error(trace)
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()
finally:
self.pool.putconn(conn)
@activity.defn(name="export_data_to_postgres") @activity.defn(name="export_data_to_postgres")
async def export_data_to_postgres(self, input_data: dict[str, Any]): async def export_data_to_postgres(self, input_data: dict[str, Any]):
@@ -141,28 +150,32 @@ class Postgres(BaseActivity):
data (DataFrame): The data to export. data (DataFrame): The data to export.
""" """
self.logger.debug(
f"Exporting data to postgres: {input_data['data']}")
schema = input_data["schema"] schema = input_data["schema"]
table_name = input_data["table_name"] table_name = input_data["table_name"]
data = DataFrame(input_data["data"]) data = DataFrame(input_data["data"])
conn = self.pool.getconn() with self.session_factory() as session:
try:
data.to_sql(table_name, self.engine, schema=schema,
if_exists="append", index=False)
session.commit()
try: except Exception as e:
data.to_sql(table_name, conn, schema=schema, trace = traceback.format_exc()
if_exists="append", index=False) self.notification_handler.build_and_send_notification(
conn.commit() 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
)
except Exception as e: self.logger.error(trace)
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: finally:
self.pool.putconn(conn) session.close()

42
laborious/utils/connectors_config.py Normal file
View File

@@ -0,0 +1,42 @@
from os import getenv
import json
def build_postgres_config():
return {
'host': getenv('POSTGRES_HOST', 'localhost'),
'port': int(getenv('POSTGRES_PORT', '5432')),
'user': getenv('POSTGRES_USER', 'sientia'),
'password': getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20'))
}
def build_mlflow_config():
return {
'host': getenv('MLFLOW_HOST', 'http://localhost'),
'port': int(getenv('MLFLOW_PORT', '5080')),
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
'password': getenv('MLFLOW_PASSWORD', 'aignosi')
}
def build_opc_config():
opc_raw = getenv('OPC_CONFIG', None)
if opc_raw:
return json.loads(opc_raw)
return {
'opc': {
'name': getenv('OPC_NAME', 'opc'),
'url': getenv('OPC_URL', 'opc.tcp://localhost:4840'),
'server_uri': getenv('OPC_SERVER_URI', 'opc.tcp://localhost:4840'),
'cert_path': getenv('OPC_CERT_PATH', None),
'private_key_path': getenv('OPC_PRIVATE_KEY_PATH', None),
'server_cert_path': getenv('OPC_SERVER_CERT_PATH', None),
'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120'))
}
}

6
laborious/utils/filters/conditional_filters.py
View File

@@ -1,13 +1,11 @@
from typing import List
from pandas import DataFrame from pandas import DataFrame
def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool: def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool:
""" """
Returns True if the data is empty, False otherwise. Returns True if the specific columns have null values, False otherwise.
""" """
return data[ return not data[
data['variable'].isin(config['VARIABLES']) & data['value'].isna()].empty data['variable'].isin(config['VARIABLES']) & data['value'].isna()].empty

30
laborious/utils/git_clone.py
View File

@@ -1,30 +0,0 @@
import os
from git import Repo
from urllib.parse import quote
# Lê variáveis de ambiente
GIT_TOKEN = os.getenv("GIT_TOKEN")
GIT_EMAIL = os.getenv("GIT_EMAIL")
REPO_URL = os.getenv("REPO_URL") # ex: "github.com/usuario/repositorio.git"
CLONE_DIR = os.getenv("CLONE_DIR", "./repo_clonado")
if not GIT_TOKEN or not GIT_EMAIL or not REPO_URL:
raise EnvironmentError("As variáveis GIT_TOKEN, GIT_EMAIL e REPO_URL devem estar definidas.")
# Escapa o token (caso contenha caracteres especiais)
safe_token = quote(GIT_TOKEN)
# Monta URL com autenticação via token
repo_url_with_auth = f"https://{safe_token}@{REPO_URL}"
# Clona o repositório
print(f"Clonando repositório em {CLONE_DIR}...")
Repo.clone_from(repo_url_with_auth, CLONE_DIR)
print("Repositório clonado com sucesso.")
# Opcional: configura o e-mail globalmente no Git (ou dentro do repo)
repo = Repo(CLONE_DIR)
with repo.config_writer() as git_config:
git_config.set_value("user", "email", GIT_EMAIL)
print(f"E-mail configurado como {GIT_EMAIL}.")

22
laborious/utils/logger.py Normal file
View File

@@ -0,0 +1,22 @@
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

9
laborious/utils/policies.py Normal file
View File

@@ -0,0 +1,9 @@
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
)

144
laborious/utils/repository/opc_repository.py
View File

@@ -3,20 +3,39 @@ from asyncua.sync import Client
from asyncua.crypto.security_policies import SecurityPolicyBasic256 from asyncua.crypto.security_policies import SecurityPolicyBasic256
from asyncua.ua import DataValue, Variant, VariantType from asyncua.ua import DataValue, Variant, VariantType
from logging import Logger 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 = { data_type_map = {
'float': VariantType.Float, 'float': {
'double': VariantType.Double, 'converter': float,
'int': VariantType.Int32, 'opc_type': VariantType.Float,
'bool': VariantType.Boolean, },
'str': VariantType.String, 'double': {
'datetime': VariantType.DateTime, '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(): class OpcRepository():
def __init__(self, name: str, url: str, logger: Logger, server_uri: str, def __init__(self, name: str, url: str, logger: Logger, notification_handler: NotificationHandler,
cert_path: str = None, private_key_path: str = None, server_cert_path: str = None): 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.url = url
self.name = name self.name = name
self.server_uri = server_uri self.server_uri = server_uri
@@ -24,8 +43,10 @@ class OpcRepository():
self.private_key_path = private_key_path self.private_key_path = private_key_path
self.server_cert_path = server_cert_path self.server_cert_path = server_cert_path
self.logger = logger self.logger = logger
self.non_receive_count = 0 self.error_count = 0
self.reconnection_interval = reconnection_interval
self.last_reconnection_time = None
self.notification_handler = notification_handler
self.client = None self.client = None
def set_security(self): def set_security(self):
@@ -67,13 +88,6 @@ class OpcRepository():
self.client.secure_channel_timeout = 10000000 self.client.secure_channel_timeout = 10000000
self.client.session_timeout = 10000000 self.client.session_timeout = 10000000
def connect(self):
self.client = Client(self.url)
if self.security:
self.set_security()
self.logger.info('Starting connection...')
self.client.connect()
def connect(self): def connect(self):
""" """
Establishes a connection to the OPC server. Establishes a connection to the OPC server.
@@ -88,7 +102,24 @@ class OpcRepository():
if self.cert_path: if self.cert_path:
self.set_security() self.set_security()
self.logger.info('Starting connection...') self.logger.info('Starting connection...')
self.client.connect() 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): def disconnect(self):
self.client.disconnect() self.client.disconnect()
@@ -98,9 +129,74 @@ class OpcRepository():
def __del__(self): def __del__(self):
self.disconnect() self.disconnect()
def write_data(self, node, value, data_type, logger): def validate_connection(self):
node = self.client.get_node(node) if self.client is None:
data = float(value) return self.connect()
logger.info(f'Writing {data} - {type(data)} to {node}')
ua_data = DataValue(Variant(data, data_type_map[data_type])) if self.error_count > 5:
node.write_value(ua_data) 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

87
laborious/worker/worker.py
View File

@@ -2,30 +2,32 @@ from temporalio import workflow, client
from temporalio.worker import Worker from temporalio.worker import Worker
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.activities.activities import Activities
import os import os
import logging import asyncio
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
from laborious.workflows.sub_workflows.format_and_export_prediction import \
FormatAndExportPrediction
from laborious.activities.activities import Activities
from laborious.utils.logger import get_logger
from laborious.utils.connectors_config import (
build_postgres_config,
build_mlflow_config,
build_opc_config
)
from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.handlers import NotificationHandler
async def main(): async def main():
host = os.getenv('TEMPORAL_HOST', 'localhost:7233') host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
logger = logging.getLogger(__name__) logger = get_logger(__name__)
stream_handler = logging.StreamHandler()
stream_handler.setLevel(
os.getenv('LOG_LEVEL', 'INFO').upper()
)
stream_handler.setFormatter(
logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
)
logger.addHandler(stream_handler) logger.info('Starting Worker...')
logger.info('Starting Notification Handler...')
notification_handler = NotificationHandler( notification_handler = NotificationHandler(
servers=os.getenv('NOTIFICATION_SERVERS', 'http://localhost:29092'), servers=os.getenv('KAFKA_SERVERS', 'http://localhost:9092'),
logger=logger, logger=logger,
project_name=os.getenv('PROJECT_NAME', 'laborious'), project_name=os.getenv('PROJECT_NAME', 'laborious'),
pipeline_name='-', pipeline_name='-',
@@ -34,46 +36,31 @@ async def main():
model='-' model='-'
) )
postgres_config = { logger.info('Starting Activities...')
'host': os.getenv('POSTGRES_HOST', 'localhost'),
'port': int(os.getenv('POSTGRES_PORT', '5432')),
'user': os.getenv('POSTGRES_USER', 'sientia'),
'password': os.getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': os.getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(os.getenv('POSTGRES_MIN_CONNECTIONS', '5')),
'max_connections': int(os.getenv('POSTGRES_MAX_CONNECTIONS', '20'))
}
mlflow_config = {
'host': os.getenv('MLFLOW_HOST', 'localhost'),
'port': int(os.getenv('MLFLOW_PORT', '5000')),
'username': os.getenv('MLFLOW_USERNAME', 'aignosi'),
'password': os.getenv('MLFLOW_PASSWORD', 'aignosi')
}
opc_config = {
'name': os.getenv('OPC_NAME', 'opc'),
'url': os.getenv('OPC_URL', 'opc.tcp://localhost:4840'),
'server_uri': os.getenv('OPC_SERVER_URI', 'opc.tcp://localhost:4840'),
'cert_path': os.getenv('OPC_CERT_PATH', None),
'private_key_path': os.getenv('OPC_PRIVATE_KEY_PATH', None),
'server_cert_path': os.getenv('OPC_SERVER_CERT_PATH', None)
}
activities = Activities( activities = Activities(
postgres_config=postgres_config, postgres_config=build_postgres_config(),
mlflow_config=mlflow_config, mlflow_config=build_mlflow_config(),
opc_config=opc_config, opc_config=build_opc_config(),
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler
) )
temporal_client = await client.Client.connect(target_host=host) logger.info('Starting Temporal Client...')
temporal_client = await client.Client.connect(
target_host=host,
namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')
)
logger.info('Starting Workers...')
workers = [ workers = [
Worker( Worker(
temporal_client, temporal_client,
task_queue='predictions', task_queue='predictions-queue',
workflows=[PredictionsBatch], workflows=[PredictionsBatch, PredictionProcess,
FormatAndExportPrediction],
activities=[ activities=[
# Base # Base
activities.prepare_activity, activities.prepare_activity,
@@ -97,9 +84,13 @@ async def main():
) )
] ]
handlers = []
for w in workers: for w in workers:
await w.run() handlers.append(w.run())
logger.info('Workers started successfully')
await asyncio.gather(*handlers)
if __name__ == '__main__': if __name__ == '__main__':
import asyncio
asyncio.run(main()) asyncio.run(main())

70
laborious/workflows/predictions_batch.py
View File

@@ -3,29 +3,87 @@ from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities from laborious.activities.activities import Activities
from typing import Any from typing import Any
from laborious.utils.policies import retry_policy
from datetime import timedelta
@workflow.defn(name="predictions_batch") @workflow.defn(name="predictions_batch")
class PredictionsBatch(): class PredictionsBatch():
@workflow.run @workflow.run
async def run(self, input_data: dict[str, Any]): async def run(self, input_data: dict[str, Any]):
"""
This workflow runs a batch of predictions based on the input data.
await workflow.execute_activity_method( 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, Activities.prepare_activity,
{ {
'schedule_name': input_data['schedule_name'], 'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch' 'workflow_name': 'predictions_batch'
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
) )
data = await workflow.execute_activity_method( data = await workflow.execute_local_activity_method(
Activities.load_custom_query, Activities.load_custom_query,
input_data['query'] input_data['query'],
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
) )
input_data['data'] = data # 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( await workflow.execute_child_workflow(
'prediction_process', input_data) 'prediction_process', prediction_input)

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

@@ -3,6 +3,8 @@ from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities from laborious.activities.activities import Activities
from typing import Any from typing import Any
from datetime import timedelta
from laborious.utils.policies import retry_policy
@workflow.defn(name="format_and_export_prediction") @workflow.defn(name="format_and_export_prediction")
@@ -11,23 +13,25 @@ class FormatAndExportPrediction():
async def run(self, input_data: dict[str, Any]): async def run(self, input_data: dict[str, Any]):
""" """
This workflow formats and exports predictions based on path_flag: 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 is None: formats prediction
- If path_flag exists: creates default prediction with timestamp, model_id, confidence and comment 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 Finally exports formatted prediction to postgres table
Args: Args:
input_data(dict[str, Any]): The input data for the workflow. Contains the following keys: input_data(dict[str, Any]): The input data for the workflow.
- path_flag(str): The path flag to determine the type of prediction to format Contains the following keys:
- data(dict[str, Any]): The data to format path_flag(str): The path flag to determine the type of prediction to format
- prediction_confidence(float): The prediction confidence to be registered data(dict[str, Any]): The data to format
- timestamp(str): The timestamp of the prediction, synchronized with the data prediction_confidence(float): The prediction confidence to be registered
- model_id(str): The model id of the prediction timestamp(str): The timestamp of the prediction, synchronized with the data
- model_name(str): The model name of the prediction model_id(int): The model id of the prediction
- model_retention(str): The model retention of the prediction model_name(str): The model name of the prediction
- comment(str): The comment to be registered model_retention(str): The model retention of the prediction
- schema(str): The schema of the prediction comment(str): The comment to be registered
- table_name(str): The table name of the prediction schema(str): The schema of the prediction
- opc_servers(list[str]): The opc servers 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 opc_output_config(dict[str, Any]): The opc output config of the prediction
Returns: Returns:
bool: True if the workflow was successful, False otherwise. bool: True if the workflow was successful, False otherwise.
@@ -36,28 +40,34 @@ class FormatAndExportPrediction():
data = input_data['data'] data = input_data['data']
prediction_confidence = input_data['prediction_confidence'] prediction_confidence = input_data['prediction_confidence']
print(f"Input data: {input_data}")
if path_flag is None: if path_flag is None:
# proceed with formatting and exporting # proceed with formatting and exporting
prediction = await workflow.execute_activity_method( prediction = await workflow.execute_local_activity_method(
Activities.format_prediction, Activities.format_prediction,
{ {
'data': data, 'data': data,
'timestamp': input_data['timestamp'], 'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence, 'prediction_confidence': prediction_confidence,
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
) )
else: else:
# create default prediction # create default prediction
prediction = await workflow.execute_activity_method( prediction = await workflow.execute_local_activity_method(
Activities.format_default_prediction, Activities.format_default_prediction,
{ {
'timestamp': input_data['timestamp'], 'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence, 'prediction_confidence': prediction_confidence,
'comment': input_data['comment'] 'comment': input_data['comment']
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
) )
# write to postgres # write to postgres
@@ -67,17 +77,20 @@ class FormatAndExportPrediction():
'schema': input_data['schema'], 'schema': input_data['schema'],
'table_name': input_data['table_name'], 'table_name': input_data['table_name'],
'data': prediction 'data': prediction
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
) )
# write to opc # write to opc
opc_holder = workflow.execute_activity_method( opc_holder = workflow.execute_activity_method(
Activities.write_opc_data, Activities.write_opc_data,
{ {
'opc_servers': input_data['opc_servers'],
'opc_output_config': input_data['opc_output_config'], 'opc_output_config': input_data['opc_output_config'],
'data': prediction 'data': prediction
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
) )
await postgres_holder await postgres_holder

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

@@ -3,101 +3,144 @@ from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities from laborious.activities.activities import Activities
from typing import Any from typing import Any
from laborious.utils.policies import retry_policy
from datetime import timedelta
@workflow.defn(name="prediction_process") @workflow.defn(name="prediction_process")
class PredictionProcess(): class PredictionProcess():
@workflow.run @workflow.run
async def run(self, input_data: dict[str, Any]): 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'] data = input_data['data']
schema = input_data['schema'] schema = input_data['schema']
table_name = input_data['table_name'] table_name = input_data['table_name']
model = input_data['model'] model_id = input_data['model_id']
filters = input_data['filters']
model_name = input_data['model_name'] model_name = input_data['model_name']
model_retention = input_data['model_retention'] model_retention = input_data['model_retention']
last_timestamp = await workflow.execute_activity_method( last_timestamp = await workflow.execute_local_activity_method(
Activities.get_last_timestamp, Activities.get_last_timestamp,
{ {
'data': data 'data': data
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
) )
path_flag, confidence = await workflow.execute_activity_method( path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.input_gate, Activities.input_gate,
{ {
'filters': input_data['filters'], 'filters': input_data['input_filters'],
'data': data '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( if await self.path_flag_handler(
data, path_flag, confidence, schema, table_name, data, path_flag, input_data, confidence, last_timestamp, comment
model, last_timestamp, model_name, model_retention
): ):
return return
response_data = await workflow.execute_activity_method( response_data = await workflow.execute_local_activity_method(
Activities.request_transform, Activities.request_transform,
{ {
'data': data, 'data': data,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention 'model_retention': model_retention
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
) )
path_flag, confidence = await workflow.execute_activity_method( path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_response_gate, Activities.mlflow_response_gate,
{ {
'filters': filters, 'filters': input_data['mlflow_transform_filters'],
'data': response_data, 'data': response_data,
'type': 'transform' '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( if await self.path_flag_handler(
data, path_flag, confidence, schema, table_name, data, path_flag, input_data, confidence, last_timestamp, comment
model, last_timestamp, model_name, model_retention
): ):
return return
path_flag, confidence = await workflow.execute_activity_method( path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_content_gate, Activities.mlflow_content_gate,
{ {
'filters': filters, 'filters': input_data['mlflow_transform_filters'],
'data': response_data, 'data': response_data,
'type': 'transform' '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( if await self.path_flag_handler(
data, path_flag, confidence, schema, table_name, data, path_flag, input_data, confidence, last_timestamp, comment
model, last_timestamp, model_name, model_retention
): ):
return return
response_data = await workflow.execute_activity_method( response_data = await workflow.execute_local_activity_method(
Activities.request_predict, Activities.request_predict,
{ {
'data': response_data, 'data': response_data,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention 'model_retention': model_retention
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
) )
path_flag, confidence = await workflow.execute_activity_method( path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_response_gate, Activities.mlflow_response_gate,
{ {
'filters': filters, 'filters': input_data['mlflow_predict_filters'],
'data': response_data, 'data': response_data,
'type': 'predict' '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( if await self.path_flag_handler(
data, path_flag, confidence, schema, table_name, data, path_flag, input_data, confidence, last_timestamp, comment
model, last_timestamp, model_name, model_retention
): ):
return return
@@ -108,60 +151,78 @@ class PredictionProcess():
'data': response_data['content'], 'data': response_data['content'],
'prediction_confidence': confidence, 'prediction_confidence': confidence,
'timestamp': response_data['timestamp'], 'timestamp': response_data['timestamp'],
'model_id': model, 'model_id': model_id,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention 'model_retention': model_retention,
'opc_output_config': input_data['opc_output_config']
} }
) )
async def path_flag_handler(self, data: dict[str, Any], path_flag: str, async def path_flag_handler(self, data: dict[str, Any], path_flag: str,
confidence: int, schema: str, table_name: str, input_data: dict[str, Any], confidence: int,
model: str, last_timestamp: str, model_name: str, last_timestamp: str, comment: str):
model_retention: str):
""" """
This function handles the path flag and the confidence of the prediction. 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. It returns True if the prediction should be stopped. If path_flag is 'repeat',
If path_flag is 'continue', it calls the write workflow. If path_flag is 'stop', it stops the prediction process. 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: Args:
data (dict[str, Any]): The data to be used for the prediction. 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 path_flag (str): The path flag to determine the type of prediction to format
confidence (int): The confidence of the prediction confidence (int): The confidence of the prediction
schema (str): The schema of the prediction schema (str): The schema of the prediction
table_name (str): The table name of the prediction table_name (str): The table name of the prediction
model (str): The model id of the prediction model_id (int): The model id of the prediction
last_timestamp (str): The timestamp of the last prediction last_timestamp (str): The timestamp of the last prediction
model_name (str): The model name of the prediction model_name (str): The model name of the prediction
model_retention (str): The model retention of the prediction model_retention (int): The model retention of the prediction
comment (str): The comment of the prediction
Returns: Returns:
bool: True if the prediction should be stopped, False otherwise. bool: True if the prediction should be stopped, False otherwise.
""" """
if path_flag == 'stop':
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 return True
elif path_flag == 'repeat': elif path_flag == 'REPEAT':
# repeat last prediction # repeat last prediction
await workflow.execute_activity_method( await workflow.execute_activity_method(
Activities.repeat_last_prediction, Activities.repeat_last_prediction,
{ {
'schema': schema, 'schema': schema,
'table_name': table_name, 'table_name': table_name,
'model': model 'model_id': model_id
} },
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
) )
return True return True
elif path_flag == 'continue': elif path_flag == 'CONTINUE':
# call write workflow # call write workflow
workflow.execute_child_workflow( await workflow.execute_child_workflow(
'format_and_export_prediction', 'format_and_export_prediction',
{ {
'path_flag': path_flag, 'path_flag': path_flag,
'data': data, 'data': data,
'prediction_confidence': confidence, 'prediction_confidence': confidence,
'timestamp': last_timestamp, 'timestamp': last_timestamp,
'model_id': model, 'model_id': model_id,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention 'model_retention': model_retention,
'schema': schema,
'table_name': table_name,
'comment': comment,
'opc_output_config': input_data['opc_output_config']
} }
) )
return True return True

2
requirements.txt
View File

@@ -1,5 +1,7 @@
temporalio temporalio
psycopg2-binary psycopg2-binary
sqlalchemy
asyncua asyncua
redis
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git git+ssh://git@github.com/Aignosi/sientia-dataops-library.git
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git git+ssh://git@github.com/Aignosi/sientia-mlops-library.git

17
tests/laborious/activities/test_base.py
View File

@@ -1,6 +1,6 @@
from unittest.mock import MagicMock from unittest.mock import MagicMock
from laborious.activities.base import BaseActivity from laborious.activities.base import BaseActivity
from pytest import fixture from pytest import fixture, mark
from sientia_do.notifications.models import Notification from sientia_do.notifications.models import Notification
@@ -12,7 +12,8 @@ def base_activity():
) )
def test_prepare_activity(base_activity): @mark.asyncio
async def test_prepare_activity(base_activity):
base_activity.notification_handler.base_notification = Notification( base_activity.notification_handler.base_notification = Notification(
project="project", project="project",
pipeline="pipeline", pipeline="pipeline",
@@ -21,12 +22,14 @@ def test_prepare_activity(base_activity):
model_id="-", model_id="-",
) )
base_activity.prepare_activity( await base_activity.prepare_activity({
schedule_name="test_schedule", 'workflow_name': 'test_workflow',
model_name="test_model", 'schedule_name': 'test_schedule',
model_id="test_model_id", '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.schedule_name == "test_schedule"
assert base_activity.notification_handler.base_notification.model_name == "test_model" 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.model_id == "test_model_id"
assert base_activity.notification_handler.base_notification.pipeline_name == "test_workflow"

26
tests/laborious/activities/test_gates.py
View File

@@ -51,7 +51,7 @@ async def test_input_gate_specific_variables_null_values_with_stop_policy_only(
} }
result = await gates.input_gate(input_data) result = await gates.input_gate(input_data)
assert result == ('stop', -1) assert result == ('stop', -1, 'Input data with bad quality')
input_args = specific_variables_null_values_mock.call_args input_args = specific_variables_null_values_mock.call_args
assert input_args[0][0].equals(DataFrame( assert input_args[0][0].equals(DataFrame(
@@ -98,7 +98,7 @@ async def test_input_gate_specific_variables_null_values_with_continue_policy_on
} }
result = await gates.input_gate(input_data) result = await gates.input_gate(input_data)
assert result == ('continue', 2) assert result == ('continue', 2, 'Input data with bad quality')
input_args = specific_variables_null_values_mock.call_args input_args = specific_variables_null_values_mock.call_args
assert input_args[0][0].equals(DataFrame( assert input_args[0][0].equals(DataFrame(
@@ -145,7 +145,7 @@ async def test_input_gate_specific_variables_null_values_no_filtered(
} }
result = await gates.input_gate(input_data) result = await gates.input_gate(input_data)
assert result == (None, 0) assert result == (None, 0, '')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
@@ -197,7 +197,7 @@ async def test_input_gate_one_stop_policy(
} }
result = await gates.input_gate(input_data) result = await gates.input_gate(input_data)
assert result == ('stop', -1) assert result == ('stop', -1, 'Input data with bad quality')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame( assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
@@ -251,7 +251,7 @@ async def test_input_gate_one_continue_policy(
} }
result = await gates.input_gate(input_data) result = await gates.input_gate(input_data)
assert result == ('continue', 2) assert result == ('continue', 2, 'Input data with bad quality')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame( assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
@@ -305,7 +305,7 @@ async def test_input_gate_no_filtered(
} }
result = await gates.input_gate(input_data) result = await gates.input_gate(input_data)
assert result == (None, 0) assert result == (None, 0, '')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame( assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
@@ -340,7 +340,7 @@ async def test_input_gate_error(
} }
result = await gates.input_gate(input_data) result = await gates.input_gate(input_data)
assert result == (None, 0) assert result == (None, 0, '')
gates.notification_handler.build_and_send_notification.assert_called_once_with( gates.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id='INTPUT_GATE_ERROR__SPECIFIC_VARIABLES_NULL_VALUES', notification_id='INTPUT_GATE_ERROR__SPECIFIC_VARIABLES_NULL_VALUES',
@@ -389,7 +389,7 @@ async def test_mlflow_response_gate_no_filtered(
} }
result = await gates.mlflow_response_gate(input_data) result = await gates.mlflow_response_gate(input_data)
assert result == (None, 0) assert result == (None, 0, '')
api_error_filter_mock.assert_called_once_with( api_error_filter_mock.assert_called_once_with(
input_data['data'], input_data['data'],
@@ -433,7 +433,7 @@ async def test_mlflow_response_gate_filtered(
} }
result = await gates.mlflow_response_gate(input_data) result = await gates.mlflow_response_gate(input_data)
assert result == ('continue', 255) assert result == ('continue', 255, "Error")
api_error_filter_mock.assert_called_once_with( api_error_filter_mock.assert_called_once_with(
input_data['data'], input_data['data'],
@@ -480,7 +480,7 @@ async def test_mlflow_content_gate_no_filtered(
} }
result = await gates.mlflow_content_gate(input_data) result = await gates.mlflow_content_gate(input_data)
assert result == (None, 0) assert result == (None, 0, '')
nan_values_filter_mock_args = nan_values_filter_mock.call_args nan_values_filter_mock_args = nan_values_filter_mock.call_args
assert nan_values_filter_mock_args[0][0].equals(DataFrame( assert nan_values_filter_mock_args[0][0].equals(DataFrame(
@@ -504,7 +504,8 @@ async def test_mlflow_content_gate_filtered(
if x == 'path_confidence': if x == 'path_confidence':
return transform_filter_path_confidence return transform_filter_path_confidence
mlflow_content_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect mlflow_content_filter_functions_mock.__getitem__.side_effect = \
transform_filter_functions_side_effect
input_data = { input_data = {
'filters': { 'filters': {
@@ -521,7 +522,8 @@ async def test_mlflow_content_gate_filtered(
} }
result = await gates.mlflow_content_gate(input_data) result = await gates.mlflow_content_gate(input_data)
assert result == ('repeat', -1) assert result == (
'repeat', -1, "Transformed data not passed the content filter")
nan_values_filter_mock_args = nan_values_filter_mock.call_args nan_values_filter_mock_args = nan_values_filter_mock.call_args
assert nan_values_filter_mock_args[0][0].equals(DataFrame( assert nan_values_filter_mock_args[0][0].equals(DataFrame(

3
tests/laborious/activities/test_mlflow.py
View File

@@ -61,7 +61,7 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow):
mlflow.model_monitoring_repository.transform.return_value = expected_response mlflow.model_monitoring_repository.transform.return_value = expected_response
# Call the method # Call the method
response_data, timestamp = await mlflow.request_transform(input_data) response_data = await mlflow.request_transform(input_data)
# Verify the data was correctly transformed # Verify the data was correctly transformed
mock_dataframe.assert_called_once_with(input_data['data']) mock_dataframe.assert_called_once_with(input_data['data'])
@@ -75,7 +75,6 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow):
# Verify the response # Verify the response
assert response_data == expected_response assert response_data == expected_response
assert timestamp == '2024-01-02'
# Verify the repository was called with correct arguments # Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.transform.assert_called_once_with( mlflow.model_monitoring_repository.transform.assert_called_once_with(

227
tests/laborious/activities/test_opc.py
View File

@@ -1,111 +1,120 @@
from unittest.mock import patch, MagicMock from unittest.mock import patch, MagicMock, ANY, call
from pytest import fixture, mark from pytest import fixture, mark
from laborious.activities.opc import NotificationLevel
from laborious.activities.opc import OPC from laborious.activities.opc import OPC
from sientia_do.notifications.models import NotificationLevel
from unittest.mock import ANY
@patch("laborious.activities.opc.OpcRepository") @patch("laborious.activities.opc.OpcRepository")
def test___init__(mock_opc_repository): 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 = OPC(
name="test", opc_servers=servers,
url="http://localhost:8080", logger=mock_logger,
server_uri="opc.tcp://localhost:4840", notification_handler=mock_notification_handler
cert_path="",
private_key_path="",
server_cert_path="",
logger=MagicMock(),
notification_handler=MagicMock()
) )
assert opc.name == "test" assert opc.opc_servers == servers
assert opc.url == "http://localhost:8080" assert opc.logger == mock_logger
assert opc.server_uri == "opc.tcp://localhost:4840" assert opc.notification_handler == mock_notification_handler
assert opc.cert_path == "" assert opc.opc_repository['server1'] == server1
assert opc.private_key_path == "" assert opc.opc_repository['server2'] == server2
assert opc.server_cert_path == ""
assert opc.opc_repository == mock_opc_repository.return_value
mock_opc_repository.assert_called_once_with( mock_opc_repository.assert_has_calls([
name="test", call(
url="http://localhost:8080", name="server1",
server_uri="opc.tcp://localhost:4840", url="http://localhost:8080",
cert_path="", logger=mock_logger,
private_key_path="", server_uri="opc.tcp://localhost:4840",
server_cert_path="", cert_path="",
logger=opc.logger, 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,
)
])
opc.opc_repository.connect.assert_called_once() server1.connect.assert_called_once()
server2.connect.assert_called_once()
@fixture @fixture
@patch("laborious.activities.opc.OpcRepository") @patch("laborious.activities.opc.OpcRepository")
def opc(mock_opc_repository): 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( return OPC(
name="test", opc_servers=servers,
url="http://localhost:8080",
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock()
) )
@mark.asyncio WRITE_DATA_CASES = [
async def test_write_opc_data_success(opc): ('tag1', 'int', 50),
# Arrange ('tag2', 'float', 50.5),
input_data = { ('tag3', 'bool', True),
'data': { ('tag4', 'string', 'test'),
'prediction': [0.75], ]
'prediction_confidence': [0.95]
},
'opc_servers': ['server1'],
'opc_output_config': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
}
}
# Act
await opc.write_opc_data(input_data)
# Assert
opc.opc_repository.write_data.assert_any_call('tag1', 0.75, 'float')
opc.opc_repository.write_data.assert_any_call('tag2', 0.95, 'float')
assert opc.opc_repository.write_data.call_count == 2
@mark.asyncio @mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
async def test_write_opc_data_prediction_error(opc): def test_write_data_success(opc, tag, data_type, data):
# Arrange opc.write_data(server='server1', tag=tag, data=data,
input_data = { data_type=data_type, tag_type='prediction')
'data': { opc.opc_repository['server1'].write_data.assert_called_once_with(
'prediction': [0.75], tag, data, data_type)
'prediction_confidence': [0.95]
},
'opc_servers': ['server1'],
'opc_output_config': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
}
}
}
opc.opc_repository.write_data.side_effect = Exception("Test error")
# Act def test_write_data_exception(opc):
await opc.write_opc_data(input_data) opc.opc_repository['server1'].write_data.side_effect = Exception(
"Test error")
# Assert opc.write_data(server='server1', tag='tag1', data=50,
opc.notification_handler.build_and_send_notification.assert_called_with( data_type='int', tag_type='prediction')
opc.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="WRITE_OPC_PREDICTION_ERROR", notification_id="WRITE_OPC_PREDICTION_ERROR",
message="Error writing data to OPC server: Test error", message="Error writing data to OPC server: Test error",
block="write_opc_data", block="write_opc_data",
@@ -116,44 +125,48 @@ async def test_write_opc_data_prediction_error(opc):
@mark.asyncio @mark.asyncio
async def test_write_opc_data_confidence_error(opc): async def test_write_opc_data_success(opc):
# Arrange # Arrange
input_data = { input_data = {
'data': { 'data': {
'prediction': [0.75], 'prediction': [0.75],
'prediction_confidence': [0.95] 'prediction_confidence': [0.95]
}, },
'opc_servers': ['server1'],
'opc_output_config': { 'opc_output_config': {
'prediction_tags': { 'server1': {
'tag1': {'data_type': 'float'} 'prediction_tags': {
}, 'tag1': {'data_type': 'float'}
'confidence_tags': { },
'tag2': {'data_type': 'float'} 'confidence_tags': {
'tag2': {'data_type': 'float'}
}
} }
} }
} }
# Make first call succeed but second fail
def side_effect(*args, **kwargs):
if args[0] == 'tag2':
raise ValueError("Test error")
return None
opc.opc_repository.write_data.side_effect = side_effect
# Act # Act
opc.write_data = MagicMock()
await opc.write_opc_data(input_data) await opc.write_opc_data(input_data)
# Assert # Assert
opc.notification_handler.build_and_send_notification.assert_called_with( opc.write_data.assert_has_calls([
notification_id="WRITE_OPC_CONFIDENCE_ERROR", call(
message="Error writing data to OPC server: Test error", server='server1',
block="write_opc_data", tag='tag1',
level=NotificationLevel.ERROR, data=0.75,
attachment_content=ANY data_type='float',
) tag_type='prediction'
opc.logger.error.assert_called_once() )])
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 @mark.asyncio
@@ -175,4 +188,4 @@ async def test_write_opc_data_empty_config(opc):
await opc.write_opc_data(input_data) await opc.write_opc_data(input_data)
# Assert # Assert
opc.opc_repository.write_data.assert_not_called() opc.opc_repository['server1'].write_data.assert_not_called()

106
tests/laborious/utils/filters/repository/test_opc_repository.py
View File

@@ -1,106 +0,0 @@
from unittest.mock import Mock, patch, MagicMock
from pathlib import Path
from asyncua.sync import Client
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from asyncua.ua import DataValue, Variant, VariantType
from pytest import fixture
from laborious.utils.repository.opc_repository import OpcRepository
@fixture
def mock_logger():
return Mock()
@fixture
def opc_repository(mock_logger):
return OpcRepository(
name="test_repo",
url="opc.tcp://localhost:4840",
logger=mock_logger,
server_uri="urn:test:server",
cert_path="/path/to/cert.pem",
private_key_path="/path/to/key.pem",
server_cert_path="/path/to/server_cert.pem"
)
@fixture
def mock_client():
with patch('laborious.utils.repository.opc_repository.Client') as mock:
client_instance = MagicMock()
mock.return_value = client_instance
yield client_instance
def test_init(opc_repository):
assert opc_repository.name == "test_repo"
assert opc_repository.url == "opc.tcp://localhost:4840"
assert opc_repository.server_uri == "urn:test:server"
assert opc_repository.cert_path == "/path/to/cert.pem"
assert opc_repository.private_key_path == "/path/to/key.pem"
assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
assert opc_repository.non_receive_count == 0
assert opc_repository.client is None
def test_set_security(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.set_security()
mock_client.application_uri = "urn:test:server"
mock_client.set_security.assert_called_once_with(
SecurityPolicyBasic256,
certificate="/path/to/cert.pem",
private_key="/path/to/key.pem",
server_certificate="/path/to/server_cert.pem"
)
assert mock_client.secure_channel_timeout == 10000000
assert mock_client.session_timeout == 10000000
def test_set_security_missing_certificates(opc_repository):
opc_repository.cert_path = None
opc_repository.private_key_path = None
try:
opc_repository.set_security()
except ValueError as e:
assert str(
e) == "Certificate and private key paths must be provided for secure connection."
def test_connect_with_security(opc_repository, mock_client):
opc_repository.connect()
mock_client.connect.assert_called_once()
assert opc_repository.client == mock_client
def test_connect_without_security(opc_repository, mock_client):
opc_repository.cert_path = None
opc_repository.connect()
mock_client.connect.assert_called_once()
assert opc_repository.client == mock_client
def test_disconnect(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.disconnect()
mock_client.disconnect.assert_called_once()
assert opc_repository.client is None
def test_write_data(opc_repository, mock_client, mock_logger):
opc_repository.client = mock_client
mock_node = MagicMock()
mock_client.get_node.return_value = mock_node
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float", mock_logger)
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_node.write_value.assert_called_once()
mock_logger.info.assert_called_once_with(
"Writing 42.0 - <class 'float'> to " + str(mock_node))

6
tests/laborious/utils/filters/test_conditional_filters.py
View File

@@ -7,18 +7,18 @@ def test_filter_specific_variables_null_values():
assert filter_specific_variables_null_values( assert filter_specific_variables_null_values(
DataFrame( DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'VARIABLES': ['variable2']}) == True config={'VARIABLES': ['variable2']}) is False
def test_filter_specific_variables_null_values_with_null_values(): def test_filter_specific_variables_null_values_with_null_values():
assert filter_specific_variables_null_values( assert filter_specific_variables_null_values(
DataFrame( DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, None]}), {'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'VARIABLES': ['variable2']}) == False config={'VARIABLES': ['variable2']}) is True
def test_filter_empty_data(): def test_filter_empty_data():
assert filter_empty_data(DataFrame(), {}) == True assert filter_empty_data(DataFrame(), {}) is True
def test_filter_empty_data_with_data(): def test_filter_empty_data_with_data():

0
tests/laborious/utils/filters/repository/test_model_repository.py → tests/laborious/utils/repository/test_model_repository.py
View File

242
tests/laborious/utils/repository/test_opc_repository.py Normal file
View File

@@ -0,0 +1,242 @@
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_write_data_validate_connection_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=False)
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_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 - <class 'float'> 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

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

@@ -1,4 +1,4 @@
from unittest.mock import call, patch, AsyncMock from unittest.mock import call, patch, AsyncMock, ANY
from pytest import mark, fixture from pytest import mark, fixture
from laborious.activities.activities import Activities from laborious.activities.activities import Activities
@@ -28,7 +28,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
await format_and_export_prediction.run(input_data) await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call( call(
Activities.format_prediction, Activities.format_prediction,
{ {
@@ -36,7 +36,9 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
'timestamp': input_data['timestamp'], 'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'] 'prediction_confidence': input_data['prediction_confidence']
} },
retry_policy=ANY,
start_to_close_timeout=ANY
)]) )])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls([
call( call(
@@ -44,8 +46,10 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
{ {
'schema': input_data['schema'], 'schema': input_data['schema'],
'table_name': input_data['table_name'], 'table_name': input_data['table_name'],
'data': workflow_mock.execute_activity_method.return_value '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([ workflow_mock.execute_activity_method.assert_has_calls([
call( call(
@@ -53,12 +57,15 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
{ {
'opc_servers': input_data['opc_servers'], 'opc_servers': input_data['opc_servers'],
'opc_output_config': input_data['opc_output_config'], 'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_activity_method.return_value '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 == 3 assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio @mark.asyncio
@@ -80,7 +87,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
await format_and_export_prediction.run(input_data) await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call( call(
Activities.format_default_prediction, Activities.format_default_prediction,
{ {
@@ -88,7 +95,9 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'], 'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment'] 'comment': input_data['comment']
} },
retry_policy=ANY,
start_to_close_timeout=ANY
) )
]) ])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls([
@@ -97,8 +106,10 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
{ {
'schema': input_data['schema'], 'schema': input_data['schema'],
'table_name': input_data['table_name'], 'table_name': input_data['table_name'],
'data': workflow_mock.execute_activity_method.return_value '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([ workflow_mock.execute_activity_method.assert_has_calls([
@@ -107,9 +118,12 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
{ {
'opc_servers': input_data['opc_servers'], 'opc_servers': input_data['opc_servers'],
'opc_output_config': input_data['opc_output_config'], 'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_activity_method.return_value '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 == 3 assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1

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

@@ -1,4 +1,4 @@
from unittest.mock import AsyncMock, patch, call from unittest.mock import AsyncMock, patch, call, ANY
from pytest import fixture, mark from pytest import fixture, mark
from laborious.activities.activities import Activities from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
@@ -18,65 +18,78 @@ async def test_run(workflow_mock, prediction_process):
'data': {'test': 'data'}, 'data': {'test': 'data'},
'schema': 'test_schema', 'schema': 'test_schema',
'table_name': 'test_table', 'table_name': 'test_table',
'model': 'test_model', 'model_id': 1,
'filters': {'test': 'filter'}, 'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30' 'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'},
} }
# Mock the activity responses # Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [ workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp '2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate ('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform) # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_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 {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95), # mlflow_response_gate (predict) # mlflow_response_gate (predict)
('continue', 0.95, "Error"),
] ]
# Act # Act
await prediction_process.run(input_data) await prediction_process.run(input_data)
# Assert # Assert
assert workflow_mock.execute_activity_method.call_count == 7 assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})]) workflow_mock.execute_local_activity_method.assert_has_calls([
workflow_mock.execute_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, { call(Activities.input_gate, {
'filters': input_data['filters'], 'filters': input_data['input_filters'],
'data': input_data['data'] 'data': input_data['data'],
})]) 'path_priority': input_data['path_priority']
workflow_mock.execute_activity_method.assert_has_calls([ }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, { call(Activities.request_transform, {
'data': input_data['data'], 'data': input_data['data'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'] 'model_retention': input_data['model_retention']
})]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, { call(Activities.mlflow_response_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform' 'type': 'transform',
})]) 'path_priority': input_data['path_priority']
workflow_mock.execute_activity_method.assert_has_calls([ }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, { call(Activities.mlflow_content_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform' 'type': 'transform',
})]) 'path_priority': input_data['path_priority']
workflow_mock.execute_activity_method.assert_has_calls([ }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, { call(Activities.request_predict, {
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'] 'model_retention': input_data['model_retention']
})]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, { call(Activities.mlflow_response_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict' '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( workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction', 'format_and_export_prediction',
@@ -85,9 +98,10 @@ async def test_run(workflow_mock, prediction_process):
'data': 'predicted_data', 'data': 'predicted_data',
'prediction_confidence': 0.95, 'prediction_confidence': 0.95,
'timestamp': '2024-01-01', 'timestamp': '2024-01-01',
'model_id': 'test_model', 'model_id': 1,
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30' 'model_retention': '30',
'opc_output_config': input_data['opc_output_config']
} }
) )
@@ -101,27 +115,34 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
'data': {'test': 'data'}, 'data': {'test': 'data'},
'schema': 'test_schema', 'schema': 'test_schema',
'table_name': 'test_table', 'table_name': 'test_table',
'model': 'test_model', 'model_id': 1,
'filters': {'test': 'filter'}, 'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30' 'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
} }
# Mock the activity responses # Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [ workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp '2024-01-01', # get_last_timestamp
('stop', 0.95), # input_gate ('stop', 0.95, "Input data with bad quality"), # input_gate
] ]
# Act # Act
await prediction_process.run(input_data) await prediction_process.run(input_data)
# Assert # Assert
assert workflow_mock.execute_activity_method.call_count == 2 assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']}), call(Activities.get_last_timestamp, {
'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, { call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']}) '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() workflow_mock.execute_child_workflow.assert_not_called()
@@ -135,42 +156,53 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'data': {'test': 'data'}, 'data': {'test': 'data'},
'schema': 'test_schema', 'schema': 'test_schema',
'table_name': 'test_table', 'table_name': 'test_table',
'model': 'test_model', 'model_id': 1,
'filters': {'test': 'filter'}, 'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30' 'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
} }
# Mock the activity responses # Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [ workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp '2024-01-01', # get_last_timestamp
('repeat', 0.95), # input_gate ('repeat', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform) ('continue', 0.95, "Error"), # mlflow_response_gate (transform)
] ]
# Act # Act
await prediction_process.run(input_data) await prediction_process.run(input_data)
# Assert # Assert
assert workflow_mock.execute_activity_method.call_count == 4 assert workflow_mock.execute_local_activity_method.call_count == 4
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})]) call(Activities.get_last_timestamp, {'data': input_data['data']},
workflow_mock.execute_activity_method.assert_has_calls([ retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, { call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})]) 'filters': input_data['input_filters'],
workflow_mock.execute_activity_method.assert_has_calls([ '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, { call(Activities.request_transform, {
'data': input_data['data'], 'data': input_data['data'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'] 'model_retention': input_data['model_retention']},
})]) retry_policy=ANY, start_to_close_timeout=ANY)
workflow_mock.execute_activity_method.assert_has_calls([ ])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, { call(Activities.mlflow_response_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform' 'type': 'transform',
})]) 'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called()
@@ -184,49 +216,61 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': {'test': 'data'}, 'data': {'test': 'data'},
'schema': 'test_schema', 'schema': 'test_schema',
'table_name': 'test_table', 'table_name': 'test_table',
'model': 'test_model', 'model_id': 1,
'filters': {'test': 'filter'}, 'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30' 'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
} }
# Mock the activity responses # Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [ workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp '2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate ('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform) ('continue', 0.95, "Error"), # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_gate (transform) # mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
] ]
# Act # Act
await prediction_process.run(input_data) await prediction_process.run(input_data)
# Assert # Assert
assert workflow_mock.execute_activity_method.call_count == 5 assert workflow_mock.execute_local_activity_method.call_count == 5
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})]) workflow_mock.execute_local_activity_method.assert_has_calls([
workflow_mock.execute_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, { call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})]) 'filters': input_data['input_filters'],
workflow_mock.execute_activity_method.assert_has_calls([ '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, { call(Activities.request_transform, {
'data': input_data['data'], 'data': input_data['data'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'] 'model_retention': input_data['model_retention']
})]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, { call(Activities.mlflow_response_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform' 'type': 'transform',
})]) 'path_priority': input_data['path_priority']
workflow_mock.execute_activity_method.assert_has_calls([ }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, { call(Activities.mlflow_content_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform' 'type': 'transform',
})]) 'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called()
@@ -240,63 +284,75 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': {'test': 'data'}, 'data': {'test': 'data'},
'schema': 'test_schema', 'schema': 'test_schema',
'table_name': 'test_table', 'table_name': 'test_table',
'model': 'test_model', 'model_id': 1,
'filters': {'test': 'filter'}, 'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30' 'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
} }
# Mock the activity responses # Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [ workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp '2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate ('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform) ('continue', 0.95, "Error"), # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_gate (transform) # mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95), # mlflow_response_gate (predict) ('continue', 0.95, "Error"), # mlflow_response_gate (predict)
] ]
# Act # Act
await prediction_process.run(input_data) await prediction_process.run(input_data)
# Assert # Assert
assert workflow_mock.execute_activity_method.call_count == 7 assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})]) call(Activities.get_last_timestamp, {'data': input_data['data']},
workflow_mock.execute_activity_method.assert_has_calls([ retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, { call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})]) 'filters': input_data['input_filters'],
workflow_mock.execute_activity_method.assert_has_calls([ '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, { call(Activities.request_transform, {
'data': input_data['data'], 'data': input_data['data'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'] 'model_retention': input_data['model_retention']
})]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, { call(Activities.mlflow_response_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform' 'type': 'transform',
})]) 'path_priority': input_data['path_priority']
workflow_mock.execute_activity_method.assert_has_calls([ }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, { call(Activities.mlflow_content_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform' 'type': 'transform',
})]) 'path_priority': input_data['path_priority']
workflow_mock.execute_activity_method.assert_has_calls([ }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, { call(Activities.request_predict, {
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'] 'model_retention': input_data['model_retention']
})]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, { call(Activities.mlflow_response_gate, {
'filters': input_data['filters'], 'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict' 'type': 'predict',
})]) 'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called()
@@ -305,7 +361,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
async def test_path_flag_handler_stop(workflow_mock, prediction_process): async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Arrange # Arrange
data = {'test': 'data'} data = {'test': 'data'}
path_flag = 'stop' path_flag = 'STOP'
confidence = 0.95 confidence = 0.95
schema = 'test_schema' schema = 'test_schema'
table_name = 'test_table' table_name = 'test_table'
@@ -317,12 +373,12 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Act # Act
result = await prediction_process.path_flag_handler( result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name, data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention model, last_timestamp, model_name, model_retention, ""
) )
# Assert # Assert
assert result is True assert result is True
workflow_mock.execute_activity_method.assert_not_called() workflow_mock.execute_local_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called()
@@ -343,7 +399,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Act # Act
result = await prediction_process.path_flag_handler( result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name, data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention model, last_timestamp, model_name, model_retention, ""
) )
# Assert # Assert
@@ -353,8 +409,10 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
{ {
'schema': schema, 'schema': schema,
'table_name': table_name, 'table_name': table_name,
'model': model 'model_id': model
} },
retry_policy=ANY,
start_to_close_timeout=ANY
) )
workflow_mock.execute_child_workflow.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called()
@@ -364,7 +422,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
async def test_path_flag_handler_continue(workflow_mock, prediction_process): async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Arrange # Arrange
data = {'test': 'data'} data = {'test': 'data'}
path_flag = 'continue' path_flag = 'CONTINUE'
confidence = 0.95 confidence = 0.95
schema = 'test_schema' schema = 'test_schema'
table_name = 'test_table' table_name = 'test_table'
@@ -376,7 +434,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Act # Act
result = await prediction_process.path_flag_handler( result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name, data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention model, last_timestamp, model_name, model_retention, 'Prediction Process'
) )
# Assert # Assert
@@ -391,7 +449,10 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'timestamp': last_timestamp, 'timestamp': last_timestamp,
'model_id': model, 'model_id': model,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention 'model_retention': model_retention,
'schema': schema,
'table_name': table_name,
'comment': 'Prediction Process'
} }
) )
@@ -413,7 +474,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Act # Act
result = await prediction_process.path_flag_handler( result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name, data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention model, last_timestamp, model_name, model_retention, ""
) )
# Assert # Assert

48
tests/laborious/workflows/subworkflows/test_predictions_batch.py
View File

@@ -1,48 +0,0 @@
from unittest.mock import AsyncMock, call, patch
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_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'
}
await predictions_batch.run(input_data)
workflow_mock.execute_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_mock.execute_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
input_data['query']
)
])
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', input_data)
])

68
tests/laborious/workflows/test_predictions_batch.py Normal file
View File

@@ -0,0 +1,68 @@
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', {}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {}),
'model_retention': input_data.get('model_retention', 60),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT'])
}
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', prediction_input)
])