SIENTIAPDE-1030

Add unit tests for orchestrator activities and workflows

- Implement tests for Activities class, covering initialization and prepare_activity method.
- Create tests for Couchbase class, including successful and failed query loading.
- Add tests for SlotManager class, verifying OPC slot loading and active ingestor retrieval.
- Develop tests for TemporalManager class, focusing on schedule loading functionality.
- Introduce tests for Orchestrator class, ensuring proper execution of workflow activities.
- Establish a new test suite for orchestrator activities and workflows in the tests directory.
This commit is contained in:
vitor-aignosi
2025-05-30 16:24:52 -03:00
parent f6584314b2
commit 4f37b2be02
59 changed files with 1827 additions and 4377 deletions

2
.gitignore vendored
View File

@@ -13,6 +13,8 @@ docker-compose.override.yml
scouter/.file_versions/ scouter/.file_versions/
scouter/pipelines/**/triggers.yaml scouter/pipelines/**/triggers.yaml
**/postgres_data/** **/postgres_data/**
**/couchbase_data/**
**/redis_data/**
# Ignorar arquivos e diretórios de cache do Python # Ignorar arquivos e diretórios de cache do Python
__pycache__/ __pycache__/
*.pyc *.pyc

View File

@@ -15,63 +15,53 @@ services:
networks: networks:
- sientia-network - sientia-network
zookeeper: couchbase:
image: confluentinc/cp-zookeeper:7.5.1 image: couchbase/server:7.2.0
container_name: zookeeper container_name: couchbase
environment:
ZOOKEEPER_CLIENT_PORT: 2181
ZOOKEEPER_TICK_TIME: 2000
ports: ports:
- "2181:2181" - "8091:8091" # Admin UI and REST API
- "8092:8092" # Query Service (N1QL)
- "8093:8093" # Index Service
- "8094:8094" # Search Service
- "11210:11210" # Data Service (KV)
- "18091:18091" # Analytics Service (if enabled)
environment:
CB_CLUSTER_USERNAME: sientia
CB_CLUSTER_PASSWORD: sientia
CB_CLUSTER_RAMSIZE: 256
CB_CLUSTER_INDEX_RAMSIZE: 256
volumes:
- ./couchbase_data:/opt/couchbase/var
networks:
- sientia-network
healthcheck:
test: ["CMD-SHELL", "curl -f http://localhost:8091/pools/default || exit 1"]
interval: 10s
timeout: 10s
retries: 5
redis:
image: redis:7-alpine # Using a lightweight Redis image
container_name: redis
ports:
- "6379:6379"
volumes:
- ./redis_data:/data # Persist Redis data
networks: networks:
- sientia-network - sientia-network
kafka: redis-commander:
image: confluentinc/cp-kafka:7.5.1 image: rediscommander/redis-commander:latest
container_name: kafka container_name: redis-commander
environment:
REDIS_HOSTS: local:redis:6379 # Connects to the 'redis' service within the Docker network
ports:
- "8081:8081" # Access the Redis Commander UI on this port
depends_on: depends_on:
- zookeeper - redis # Ensures Redis starts before Redis Commander
ports:
- "9092:9092"
- "29092:29092"
environment:
KAFKA_BROKER_ID: 1
KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181
KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://kafka:29092,PLAINTEXT_HOST://localhost:9092
KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: PLAINTEXT:PLAINTEXT,PLAINTEXT_HOST:PLAINTEXT
KAFKA_INTER_BROKER_LISTENER_NAME: PLAINTEXT
KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1
networks: networks:
- sientia-network - sientia-network
kafka-ui:
image: provectuslabs/kafka-ui:latest
container_name: kafka-ui
ports:
- "8080:8080"
environment:
KAFKA_CLUSTERS_0_NAME: local
KAFKA_CLUSTERS_0_BOOTSTRAPSERVERS: kafka:29092
networks:
- sientia-network
simulator:
build:
context: .
dockerfile: simulator/Dockerfile
args:
GIT_REPO: ${SIMULATOR_GIT_REPO}
GIT_BRANCH: ${SIMULATOR_GIT_BRANCH}
container_name: simulator
ports:
- "4840:4840"
depends_on:
- kafka
networks:
- sientia-network
env_file:
- .env
networks: networks:
sientia-network: sientia-network:
@@ -80,3 +70,7 @@ networks:
volumes: volumes:
postgres_data: postgres_data:
driver: local driver: local
couchbase_data:
driver: local
redis_data:
driver: local

View File

@@ -1,53 +0,0 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.postgres import Postgres
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.opc import OPC
from typing import Any
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
class Activities(Postgres, MLFlow, Gates, OPC):
def __init__(self,
postgres_config: dict[str, Any],
mlflow_config: dict[str, Any],
opc_config: dict[str, Any],
logger: Logger, notification_handler: NotificationHandler):
# Initialize parent classes
Postgres.__init__(self, host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler)
MLFlow.__init__(self, mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler)
Gates.__init__(self, logger=logger,
notification_handler=notification_handler)
OPC.__init__(self,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler)
@activity.defn(name="prepare_activity")
async def prepare_activity(self, input_data: dict[str, Any]):
await super().prepare_activity(input_data)
def shutdown(self):
Postgres.close(self)
OPC.shutdown(self)

View File

@@ -1,26 +0,0 @@
from typing import Any
from logging import Logger
from temporalio import activity
from sientia_do.notifications.handlers import NotificationHandler
class BaseActivity:
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
self.logger = logger
self.notification_handler = notification_handler
@activity.defn(name="prepare_activity")
async def prepare_activity(self, input_data: dict[str, Any]):
"""
Prepare the activity for the notification handler.
Args:
workflow_name (str): The name of the workflow.
schedule_name (str): The name of the schedule.
model_name (str): The name of the model.
model_id (str): The id of the model.
"""
self.notification_handler.base_notification.pipeline_name = input_data['workflow_name']
self.notification_handler.base_notification.schedule_name = input_data['schedule_name']
self.notification_handler.base_notification.model_name = input_data['model_name']
self.notification_handler.base_notification.model_id = input_data['model_id']

View File

@@ -1,297 +0,0 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.base import BaseActivity
from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
from typing import Any
from laborious.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values
)
from pandas import DataFrame
from datetime import datetime
input_filter_functions = {
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {
'STOP': -1,
'CONTINUE': 2,
'REPEAT': -1
}
}
mlflow_response_filter_functions = {
'API_ERROR': api_error_filter,
'path_confidence': {
'STOP': -1,
'CONTINUE': 10,
'REPEAT': -1
},
}
mlflow_content_filter_functions = {
'NAN_VALUES': nan_values_filter,
'path_confidence': {
'STOP': -1,
'CONTINUE': 18,
'REPEAT': -1
}
}
class Gates(BaseActivity):
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
BaseActivity.__init__(self, logger, notification_handler)
@activity.defn(name="input_gate")
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Filters the data based on the filters. The return value is a tuple with the first element
being the policy and the second element being the confidence status.
Args:
input_data (dict): The input data. Contains:
filters (dict): The filters to apply.
The key is the filter name and the value is the filter configuration.
data (dict[str, Any]): The data to filter.
path_priority (list[str]): The path priority.
Returns:
tuple[str | None, int, str]: (policy, confidence) based in priority
list and filter configuration and functions.
"""
self.logger.debug("Performing input gate...")
filters = input_data['filters']
data = DataFrame(input_data['data'])
path_priority = input_data['path_priority']
filter_output = []
self.logger.debug(f"Input data:\n {data}")
self.logger.debug(f"Filters: {filters}")
for fil, config in filters.items():
if fil not in input_filter_functions:
self.logger.error(f"Filter {fil} not found")
continue
try:
if input_filter_functions[fil](data, config):
self.logger.debug(
f"Data not passed the input filter {fil}:{config}")
filter_output.append(config['POLICY'])
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"INTPUT_GATE_ERROR__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="input_gate",
level=NotificationLevel.ERROR,
attachment_content=trace
)
for path_flag in path_priority:
if path_flag in filter_output:
self.logger.debug(f"Input gate result: {path_flag}")
return path_flag, input_filter_functions['path_confidence'][path_flag], \
"Input data with bad quality"
self.logger.debug("Nothing was filtered by the input gate")
return None, 0, ""
@activity.defn(name="mlflow_response_gate")
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Filters the data based on the mlflow response filters.
The return value is a tuple with the first element
being the policy and the second element being the confidence status.
Args:
input_data (dict): The input data. Contains:
filters (dict): The filter configuration to apply.
data (dict[str, Any]): The data to filter.
path_priority (list[str]): The path priority list.
type (str): The type of the gate.
Returns:
tuple[str | None, int, str]: (policy, confidence) based in priority list
and filter configuration and functions.
"""
self.logger.debug("Performing mlflow response gate...")
filters = input_data['filters']
data = input_data['data']
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self.logger.debug(f"Input data:\n {data}")
self.logger.debug(f"Filters: {filters}")
comments = []
for fil, config in filters.items():
if fil not in mlflow_response_filter_functions:
continue
try:
if mlflow_response_filter_functions[fil](data, config):
filter_output.append(config['POLICY'])
comments.append(data['content']['message'])
self.notification_handler.build_and_send_notification(
notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
message=data['content']['message'],
block="mlflow_gate",
level=NotificationLevel.WARNING,
attachment_content=data['content']['traceback']
)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="mlflow_gate",
level=NotificationLevel.ERROR,
attachment_content=trace
)
for path_flag in path_priority:
if path_flag in filter_output:
self.logger.debug(f"Mlflow response gate result: {path_flag}")
return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \
", ".join(comments)
self.logger.debug("Nothing was filtered by the mlflow response gate")
return None, 0, ""
@activity.defn(name="mlflow_content_gate")
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Filters the data based on the mlflow content filters.
The return value is a tuple with the first element
being the policy and the second element being the confidence status.
Args:
input_data (dict): The input data. Contains:
filters (dict): The filter configuration to apply.
data (dict[str, Any]): The data to filter.
path_priority (list[str]): The path priority list.
type (str): The type of the gate.
Returns:
tuple[str | None, int, str]: (policy, confidence) based in priority
list and filter configuration and functions.
"""
self.logger.debug("Performing mlflow content gate...")
filters = input_data['filters']
data = DataFrame(input_data['data'])
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self.logger.debug(f"Input data:\n {data}")
self.logger.debug(f"Filters: {filters}")
for fil, config in filters.items():
if fil not in mlflow_content_filter_functions:
continue
try:
if mlflow_content_filter_functions[fil](data, config):
filter_output.append(config['POLICY'])
self.notification_handler.build_and_send_notification(
notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
message=f"Data not passed the content filter {fil}:{config}",
block="mlflow_gate",
level=NotificationLevel.WARNING,
attachment_content=data.to_string()
)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="mlflow_gate",
level=NotificationLevel.ERROR,
attachment_content=trace
)
for path_flag in path_priority:
if path_flag in filter_output:
self.logger.debug(f"Mlflow content gate result: {path_flag}")
return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \
"Transformed data not passed the content filter"
self.logger.debug("Nothing was filtered by the mlflow content gate")
return None, 0, ""
@activity.defn(name="format_prediction")
async def format_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Formats the prediction data.
Args:
input_data (dict): The input data. Contains:
data (dict[str, Any]): The data to format.
timestamp (str): The timestamp of the data.
model_id (str): The id of the model.
prediction_confidence (float): The confidence of the prediction.
Returns:
dict: The formatted data.
"""
self.logger.debug("Formatting prediction...")
data = DataFrame(input_data['data'])
data['timestamp'] = input_data['timestamp']
data['model_id'] = input_data['model_id']
data['prediction_confidence'] = input_data['prediction_confidence']
data['prediction_status'] = 'Good'
data['comments'] = ""
data.sort_values(by='timestamp', inplace=True)
return data.to_dict()
@activity.defn(name="format_default_prediction")
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Creates and formats the default prediction data, with zero value in prediction,
and usefull information in the other fields.
Args:
input_data (dict): The input data. Contains:
timestamp (str): The timestamp of the data.
model_id (str): The id of the model.
prediction_confidence (float): The confidence of the prediction.
comment (str): The comment of the prediction.
Returns:
dict: The formatted data.
"""
self.logger.debug("Formatting default prediction...")
return DataFrame({
'prediction': [0],
'response_time': [0],
'timestamp': [input_data['timestamp']],
'model_id': [input_data['model_id']],
'prediction_confidence': [input_data['prediction_confidence']],
'prediction_status': ['Bad'],
'comments': [input_data['comment']]
}).to_dict()
@activity.defn(name="get_last_timestamp")
async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
"""
Gets the last timestamp of the data.
Args:
input_data (dict): The input data. Contains:
data (dict[str, Any]): The data to get the last timestamp from.
Returns:
str: The last timestamp of the data.
"""
data = DataFrame(input_data['data'])
if data.empty:
return datetime.now().strftime('%Y-%m-%d %H:%M:%S')
return max(data['timestamp'].values.tolist())

View File

@@ -1,91 +0,0 @@
import numpy as np
from pandas import DataFrame
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.base import BaseActivity
from laborious.utils.repository.model_repository import MLFlowRepository
from typing import Any
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
class MLFlow(BaseActivity):
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
BaseActivity.__init__(self, logger, notification_handler)
self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username
self.mlflow_password = mlflow_password
self.model_monitoring_repository = MLFlowRepository(
f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password
)
@activity.defn(name="request_transform")
async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Access MLFlow model to get the transformed data.
Args:
input_data (dict): The input data. Contains:
data (dict[str, Any]): The data to transform.
model_name (str): The name of the model.
model_retention (int): The retention of the model in minutes.
Returns:
dict[str, Any]: The transformed data.
"""
self.logger.info('Transforming data...')
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_retention = input_data['model_retention']
self.logger.debug("Raw input data:")
self.logger.debug(data)
data = data.pivot(
index='timestamp', columns='variable',
values='value')
data.fillna(np.nan, inplace=True)
data.reset_index(inplace=True)
data.columns.name = None
self.logger.debug("Processed input data:")
self.logger.debug(data)
response_data = self.model_monitoring_repository.transform(
model_name, data, model_retention)
self.logger.debug("Response data:")
self.logger.debug(response_data)
return response_data
@activity.defn(name="request_predict")
async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Access MLFlow model to get the predicted data.
Args:
input_data (dict): The input data. Contains:
data (dict[str, Any]): The data to predict.
model_name (str): The name of the model.
model_retention (int): The retention of the model.
Returns:
dict[str, Any]: The predicted data.
"""
self.logger.info('Predicting data...')
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_retention = input_data['model_retention']
self.logger.debug(data)
data.replace(np.nan, None, inplace=True)
response_data = self.model_monitoring_repository.predict(
model_name, data, model_retention)
self.logger.debug(response_data)
return response_data

View File

@@ -1,105 +0,0 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.base import BaseActivity
from laborious.utils.repository.opc_repository import OpcRepository
from typing import Any
import traceback
from pandas import DataFrame
class OPC(BaseActivity):
def __init__(self, opc_servers: dict[str, dict[str, Any]],
logger: Logger, notification_handler: NotificationHandler):
self.logger = logger
self.notification_handler = notification_handler
self.opc_servers = opc_servers
self.opc_repository = {}
for name, server in opc_servers.items():
self.opc_repository[name] = OpcRepository(
name=name,
url=server['url'],
logger=self.logger,
server_uri=server['server_uri'],
cert_path=server['cert_path'],
private_key_path=server['private_key_path'],
server_cert_path=server['server_cert_path'],
notification_handler=self.notification_handler,
reconnection_interval=server['reconnection_interval'],
)
self.opc_repository[name].connect()
BaseActivity.__init__(self, logger, notification_handler)
def write_data(self, server: str, tag: str, data: Any,
data_type: str, tag_type: str):
try:
self.opc_repository[server].write_data(
tag, data, data_type)
self.logger.debug(f"Wrote {tag_type} to {tag}")
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
message=f"Error writing data to OPC server: {e}",
block="write_opc_data",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
@activity.defn(name='write_opc_data')
async def write_opc_data(self, input_data: dict[str, Any]):
"""
Write prediction and confidence data to OPC servers. The two writing
operations are optional and independent of each other.
Args:
input_data (dict[str, Any]): The input data. Contains the following keys:
- data (dict[str, Any]): The dataframe that contains the data to write
to the OPC servers.
- opc_output_config (dict[str, Any]): The OPC writing configuration.
The keys are the OPC server names and the values contain:
prediction_tags (dict[str, Any]): The tags to write to the OPC servers.
confidence_tags (dict[str, Any]): The tags to write to the OPC servers.
Returns:
"""
self.logger.debug("Writing data to OPC servers...")
data = DataFrame(input_data['data'])
opc_output_config = input_data['opc_output_config']
self.logger.debug(data)
for server, config in opc_output_config.items():
if self.opc_repository.get(server) is None:
self.logger.error(f"OPC server {server} not found")
continue
if 'prediction_tags' in config:
for tag, tag_config in config['prediction_tags'].items():
self.write_data(
server=server,
tag=tag,
data=data.head(1)['prediction'].values[0],
data_type=tag_config['data_type'],
tag_type='prediction'
)
if 'confidence_tags' in config:
for tag, tag_config in config['confidence_tags'].items():
self.write_data(
server=server,
tag=tag,
data=data.head(1)['prediction_confidence'].values[0],
data_type=tag_config['data_type'],
tag_type='confidence'
)
def shutdown(self):
for opc in self.opc_repository.values():
opc.disconnect()

View File

@@ -1,181 +0,0 @@
import traceback
from temporalio import workflow, activity
from laborious.activities.base import BaseActivity
with workflow.unsafe.imports_passed_through():
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool
from psycopg2.pool import ThreadedConnectionPool
from pandas import read_sql_query, DataFrame
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from typing import Any
class Postgres(BaseActivity):
def __init__(self, host: str, port: int,
user: str, password: str, dbname: str,
min_connections: int, max_connections: int,
logger: Logger, notification_handler: NotificationHandler):
self.host = host
self.port = port
self.user = user
self.password = password
self.dbname = dbname
# Create SQLAlchemy engine with connection pooling
self.engine = create_engine(
f'postgresql://{user}:{password}@{host}:{port}/{dbname}',
poolclass=QueuePool,
pool_size=min_connections,
max_overflow=max_connections - min_connections,
pool_pre_ping=True
)
self.session_factory = sessionmaker(bind=self.engine)
BaseActivity.__init__(self, logger, notification_handler)
def close(self):
self.engine.dispose()
def __del__(self):
self.close()
@activity.defn(name="load_custom_query")
async def load_custom_query(self, query: str) -> dict[str, Any]:
"""
Loads data from a custom query.
Args:
query (str): The query to load data from.
Returns:
dict[str, dict]: The data from the query.
"""
self.logger.info(f"Fetching data from query: {query}")
data = None
with self.session_factory() as session:
try:
data = read_sql_query(query, self.engine)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id="ERROR_LOADING_CUSTOM_QUERY",
message=f"Error fetching data from query: {e}",
block="load_custom_query",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
return {}
finally:
session.close()
if data is None:
return {}
# Converts any datetime datatype columns to string
for col in data.select_dtypes(include=['datetime64']).columns:
data[col] = data[col].dt.strftime('%Y-%m-%d %H:%M:%S')
self.logger.info(f"Fetched {len(data)} rows")
self.logger.debug(f"Data: \n{data.to_string()}")
return data.to_dict()
@activity.defn(name="repeat_last_prediction")
async def repeat_last_prediction(self, query_items: dict[str, str]):
"""
Repeats the last prediction for a given model.
Args:
query_items (dict[str, str]): The query items. Contains:
schema (str): The schema of the table.
table_name (str): The name of the table.
model (int): The model to repeat the prediction for.
Returns:
None
"""
schema = query_items["schema"]
table_name = query_items["table_name"]
model = query_items["model"]
repeat_query = f"""
INSERT INTO \"{schema}\".{table_name} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at)
SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW()
FROM \"{schema}\".{table_name}
WHERE model_id = {model}
ORDER BY timestamp DESC
LIMIT 1;
"""
self.logger.info(f"Repeating last prediction for model {model}")
self.logger.debug(f"Query: {repeat_query}")
with self.session_factory() as session:
try:
session.execute(repeat_query)
session.commit()
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id="ERROR_REPEATING_LAST_PREDICTION",
message=f"Error repeating last prediction: {e}",
block="repeat_last_prediction",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
finally:
session.close()
@activity.defn(name="export_data_to_postgres")
async def export_data_to_postgres(self, input_data: dict[str, Any]):
"""
Exports data to a postgres table.
Args:
input_data (dict[str, Any]): The data to export. Contains:
schema (str): The schema of the table.
table_name (str): The name of the table.
data (DataFrame): The data to export.
"""
self.logger.debug(
f"Exporting data to postgres: {input_data['data']}")
schema = input_data["schema"]
table_name = input_data["table_name"]
data = DataFrame(input_data["data"])
with self.session_factory() as session:
try:
data.to_sql(table_name, self.engine, schema=schema,
if_exists="append", index=False)
session.commit()
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
message=f"Error exporting data to postgres: {e}",
block="export_data_to_postgres",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
else:
self.logger.debug("Data exported to postgres")
finally:
session.close()

View File

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

View File

@@ -1,22 +0,0 @@
import numpy as np
from pandas import DataFrame
def api_error_filter(response: dict, _config: dict):
if not response:
return True
if not response['success']:
return True
return False
def nan_values_filter(predictions: DataFrame, _config: dict):
data = predictions.replace({None: np.nan}).drop(
columns=['timestamp'], errors='ignore').infer_objects(copy=False)
if data.isna().all().all():
return True
return False

View File

@@ -1,22 +0,0 @@
from os import getenv
import logging
import sys
def get_logger(name: str):
log_level = getenv('LOG_LEVEL', 'INFO').upper()
logger = logging.getLogger(name)
logger.setLevel(log_level)
stream_handler = logging.StreamHandler(sys.stdout)
stream_handler.setLevel(log_level)
stream_handler.setFormatter(
logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
)
logger.addHandler(stream_handler)
return logger

View File

@@ -1,9 +0,0 @@
from datetime import timedelta
from temporalio.common import RetryPolicy
retry_policy = RetryPolicy(
initial_interval=timedelta(seconds=1),
backoff_coefficient=2.0,
maximum_interval=timedelta(minutes=1),
maximum_attempts=1
)

View File

@@ -1,297 +0,0 @@
"""
Model Monitoring Repository
This module contains the ModelMonitoringRepository class, which is responsible for handling the communication with the Model Monitoring API.
It includes the methods that are used to answer ModelMonitoringService requests using the Model Monitoring API functions.
By Monitoring we mean the evaluation of the performance of models, the generation of reports.
"""
from datetime import datetime
import traceback
import mlflow
import pandas as pd
from sientia.ModelServing import ModelServing
class MLFlowRepository():
def __init__(self, host, username, password):
self.model_serving = ModelServing(tracking_uri=host,
username=username, password=password)
def get_current_data_df(self, current_data: pd.DataFrame, model_name: str, target: str):
"""
Get the current data as a DataFrame and update the prediction and target columns
Parameters:
current_data (pd.DataFrame): the current data
model_name (str): the name of the model
target (str): the target column
Returns:
DataFrame: the current data as a DataFrame
"""
predictions = current_data['prediction']
target = current_data[target]
current_data = self.model_serving.get_transformed_data(
model_name, current_data, by='model')
current_data['prediction'] = predictions
current_data['target'] = target
return pd.DataFrame(current_data).dropna()
def get_artifact(self, destination: str, search_by: str, run_id: str = None,
model_name: str = None, artifact_name: str = None) -> None:
"""
Get an artifact in MLflow by experiment or model and save it to a destination path using API.
If the artifact is searched by model, the latest production version will be used.
Args:
destination: The destination path to save the artifact.
search_by: The way to search for the artifact ('experiment' or 'model').
run_id: The run ID of the experiment (if search_by is "experiment").
model_name: The name of the model (if search_by is "model").
artifact_name: The path of the artifact to download.
Returns:
artifact: The artifact(.csv) downloaded from MLflow.
"""
self.model_serving.get_artifact(destination=destination, search_by=search_by,
run_id=run_id, model_name=model_name, artifact_name=artifact_name)
def calculate_model_metrics(self, real_data, predictions, flag):
"""
Function to calculate the metrics of a model using API
Parameters:
real_data (array): the real data
predictions (array): the predictions
Returns:
dict: the metrics of the model including MSE and R2
"""
return self.model_serving.get_model_metrics(reference_data=None, real_data=real_data, predictions=predictions, type_flag=flag)
def get_experiment_by_run_id(self, run_id: str) -> dict:
# Get the run information using the run_id
run = mlflow.get_run(run_id)
# Extract the experiment ID from the run
experiment_id = run.info.experiment_id
# Get the experiment details using the experiment ID
experiment = mlflow.get_experiment(experiment_id)
experiment_name = experiment.name
return experiment_name
def get_next_run_name(self, model_name: str) -> str:
"""
Function to get the next run number of a specific model
Parameters:
model_name (str): the name of the model
Returns:
str: the next run number
"""
runs = mlflow.search_runs(
experiment_names=[model_name], order_by=["start_time desc"])
next_run_number = len(runs) + 1
return f"{model_name}-{next_run_number}"
def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
"""
Retrain a model with new data.
Parameters:
data (pandas.DataFrame): The new data to use for retraining.
model_name (str): The name of the model to retrain.
metrics_list (list): The metrics to be used to compare the models.
compare_metrics (bool): If True, the retrain will only be considered if the new model is better than the current one.
If False, the retrain will always be considered.
split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets.
If False, the data will be used as a unique block for retraining.
update_report (bool): If True, a report will be created with the data of the retrained model.
update_transformation (bool): If True, the model will be updated in the MLflow tracking server.
update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server.
shuffle_data (bool): If True, the data will be shuffled before splitting.
model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'.
Returns:
mlflow.sklearn.Model: The retrained prediction model.
mlflow.sklearn.Model: The retrained data model.
mse (float): The mean squared error of the retrained model.
r2 (float): The R-squared score of the retrained model.
"""
# load predictor model
predictor_uri = f"models:/{model_name}/production"
# load transform model
latest_production_id = self.model_serving.get_model_run_id(
model_name, stage="Production"
)
transform_uri = self.model_serving.get_model_uri(
latest_production_id, prediction=False
)
# load
data_model = mlflow.sklearn.load_model(transform_uri)
prediction_model = mlflow.sklearn.load_model(predictor_uri)
data_model = data_model.fit(data)
treated_data = data_model.predict(data)
# align target column with treated_data
target_name = data_model.target_variable
y = data[target_name]
treated_data = pd.merge(
treated_data, y, left_index=True, right_index=True)
prediction_model = prediction_model.fit(treated_data)
# Example usage
experiment = self.get_experiment_by_run_id(latest_production_id)
pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_model) # load class attributes
mlflow.set_experiment(experiment)
experiment_description = "Retrain model {model_name} with new data"
current_run_name = self.get_next_run_name(experiment)
with mlflow.start_run(
run_name=current_run_name, description=experiment_description
) as _run:
# update transfomation model
# fixed parameters
for name_atribute, val_atribute in pred_model_atributes.items():
if name_atribute != "model":
mlflow.log_param(name_atribute, val_atribute)
# update prediction model
for name_atribute, val_atribute in data_model_atributes.items():
if name_atribute != "model":
mlflow.log_param(name_atribute, val_atribute)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(data_model, "data_model")
file_path = f"laborious/data/raw_data_{model_name}.csv"
data.to_csv(
f"laborious/data/raw_data_{model_name}.csv", index=True)
# log the data raw
mlflow.log_artifact(file_path)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(prediction_model, "prediction_model")
mlflow.log_param("retrain", True)
return "Model retrained successfully", experiment
def get_experiment(self, experiment_name: str) -> int:
experiment = mlflow.get_experiment_by_name(experiment_name)
if experiment is None:
raise ValueError(f'Experiment {experiment_name} not found')
return int(experiment.experiment_id)
def get_experiment_last_run(self, experiment_id: int) -> str:
runs = mlflow.search_runs(
experiment_ids=[experiment_id],
filter_string="", # Sem filtro no MLflow ainda
output_format="pandas"
)
# Filtrar apenas as runs onde params.retrain == True
filtered_runs = runs[runs["params.retrain"] == 'True']
# Converter a coluna 'end_time' para datetime
filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])
# Ordenar o DataFrame de forma descendente pela coluna 'end_time'
filtered_runs = filtered_runs.sort_values(
by='end_time', ascending=False)
# Pegar a última run_id do DataFrame filtrado e ordenado
latest_run_id = filtered_runs.iloc[0]['run_id']
return latest_run_id
def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict:
# Registrar o modelo
# Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
# Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
mlflow.register_model(
f"runs:/{run_id}/prediction_model", model_name)
# Colocar a versão do modelo em produção
# Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production'
client = mlflow.tracking.MlflowClient()
# Obter a versão mais recente registrada do modelo
model_versions = client.get_registered_model(
model_name).latest_versions
max_version = max(model_versions, key=lambda x: int(x.version)).version
# Mover a versão mais recente do modelo para o estágio de 'Production'
client.transition_model_version_stage(
name=model_name,
version=max_version,
stage="Production",
archive_existing_versions=True
)
return {
'model_name': model_name,
'version': max_version,
'mlflow_run_id': run_id
}
def update_production_model(self, experiment: str, model_name: str) -> dict:
experiment_id = self.get_experiment(experiment)
run_id = self.get_experiment_last_run(experiment_id)
metadata = self.update_production_model_by_run_id(run_id, model_name)
metadata['mlflow_experiment_id'] = experiment_id
return metadata
def transform(self, model_name: str, data: pd.DataFrame, model_retention: int):
try:
return {
'success': True,
'content': self.model_serving.get_cached_transform(
model_name, data, model_retention).to_dict()
}
except Exception as e:
return {
'success': False,
'content': {
'message': str(e),
'traceback': traceback.format_exc()
}
}
def predict(self, model_name: str, data: pd.DataFrame, model_retention: int):
try:
start_time = datetime.now()
data = self.model_serving.get_cached_predict(
model_name, data, model_retention)[-1:]
end_time = datetime.now()
data = pd.DataFrame(data, columns=['prediction'])
data['response_time'] = (end_time - start_time).total_seconds()
return {
'success': True,
'content': data.to_dict()
}
except Exception as e:
return {
'success': False,
'content': {
'message': str(e),
'traceback': traceback.format_exc()
}
}

View File

@@ -1,207 +0,0 @@
from pathlib import Path
from asyncua.sync import Client
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from asyncua.ua import DataValue, Variant, VariantType
from logging import Logger
from datetime import datetime
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
import traceback
data_type_map = {
'float': {
'converter': float,
'opc_type': VariantType.Float,
},
'double': {
'converter': float,
'opc_type': VariantType.Double,
},
'int': {
'converter': int,
'opc_type': VariantType.Int32,
},
'bool': {
'converter': bool,
'opc_type': VariantType.Boolean,
},
'str': {
'converter': str,
'opc_type': VariantType.String,
}
}
class OpcRepository():
def __init__(self, name: str, url: str, logger: Logger, notification_handler: NotificationHandler,
reconnection_interval: int = 60, server_uri: str = None, cert_path: str = None,
private_key_path: str = None, server_cert_path: str = None):
self.url = url
self.name = name
self.server_uri = server_uri
self.cert_path = cert_path
self.private_key_path = private_key_path
self.server_cert_path = server_cert_path
self.logger = logger
self.error_count = 0
self.reconnection_interval = reconnection_interval
self.last_reconnection_time = None
self.notification_handler = notification_handler
self.client = None
def set_security(self):
"""
Configures the security settings for the OPC UA client.
This method sets up the security policy, certificates, and timeouts
required for establishing a secure connection with the OPC UA server.
Raises:
ValueError: If either the certificate path or private key path is not provided.
Attributes:
cert_path (str): Path to the client's certificate file.
private_key_path (str): Path to the client's private key file.
server_cert_path (str, optional): Path to the server's certificate file.
server_uri (str): The URI of the server to be used as the application URI.
client (opcua.Client): The OPC UA client instance.
logger (logging.Logger): Logger instance for logging information.
Security Settings:
- Security Policy: Basic256
- Secure Channel Timeout: 10,000,000 ms
- Session Timeout: 10,000,000 ms
"""
if not all([self.cert_path, self.private_key_path]):
raise ValueError(
"Certificate and private key paths must be provided for secure connection.")
cert = Path(self.cert_path)
private_key = Path(self.private_key_path)
server_cert = Path(
self.server_cert_path) if self.server_cert_path else None
self.client.application_uri = self.server_uri
self.logger.info('Setting security...')
self.client.set_security(
SecurityPolicyBasic256,
certificate=str(cert),
private_key=str(private_key),
server_certificate=str(server_cert)
)
self.client.secure_channel_timeout = 10000000
self.client.session_timeout = 10000000
def connect(self):
"""
Establishes a connection to the OPC server.
This method initializes the OPC client using the provided URL and
sets up security if a certificate path is specified. It then
attempts to connect to the server and logs the connection status.
Raises:
Exception: If the connection to the OPC server fails.
"""
self.client = Client(self.url)
if self.cert_path:
self.set_security()
self.logger.info('Starting connection...')
return self.try_connect()
def try_connect(self):
try:
self.last_reconnection_time = datetime.now()
self.client.connect()
return True
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"OPC_CONNECTION_ERROR_{self.name}",
message=f"Failed to connect to OPC server: {e}",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
return False
def disconnect(self):
if self.client is None:
return
self.client.disconnect()
self.client = None
self.logger.info('Disconnected from OPC server')
def __del__(self):
try:
self.disconnect()
except Exception as e:
self.logger.error(f"Error in destructor: {e}")
def validate_connection(self):
if self.client is None:
return self.connect()
if self.error_count > 5:
self.logger.warning(
f"OPC server {self.name} will be disconnected due to multiple errors")
try:
self.disconnect()
except Exception as e:
trace = traceback.format_exc()
self.logger.error(f"Failed to disconnect from OPC server: {e}")
self.logger.error(trace)
self.logger.info(
f"Attempting to reconnect to OPC server {self.name}...")
return self.connect()
if hasattr(self.client, 'aio_obj') and self.client.aio_obj.uaclient.protocol is None or \
(hasattr(self.client.aio_obj.uaclient, 'protocol') and
self.client.aio_obj.uaclient.protocol.state == "closed"):
self.logger.error(
f"OPC server {self.name} is not connected")
if (datetime.now() - self.last_reconnection_time).total_seconds(
) > self.reconnection_interval:
self.logger.error(
f"Trying to reconnect to OPC server {self.name}...")
return self.try_connect()
return False
return True
def write_data(self, node, value, data_type):
if not self.validate_connection():
return
try:
node = self.client.get_node(node)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"OPC_WRITE_GET_NODE_ERROR_{self.name}",
message=f"Failed to get node from OPC server: {e}",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
self.error_count += 1
return
data = data_type_map[data_type]['converter'](value)
self.logger.info(f'Writing {data} - {type(data)} to {node}')
ua_data = DataValue(
Variant(data, data_type_map[data_type]['opc_type']))
try:
node.write_value(ua_data)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"OPC_WRITE_DATA_ERROR_{self.name}",
message=f"Failed to write data to OPC server: {e}",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.logger.error(trace)
self.error_count += 1
return
self.error_count = 0

View File

@@ -1,89 +0,0 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from laborious.utils.policies import retry_policy
from datetime import timedelta
@workflow.defn(name="predictions_batch")
class PredictionsBatch():
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
This workflow runs a batch of predictions based on the input data.
The workflow executes in two main steps:
1. Prepares the activity with schedule and model information
2. Loads data using a custom query and executes the prediction process
Args:
input_data (dict[str, Any]): The input data for the workflow.
Contains the following keys:
schedule_name (str): The name of the schedule.
model_name (str): The name of the model.
model_id (int): The id of the model.
query (str): The SQL query to be executed to load data.
schema (dict, optional): The schema definition for the data.
table_name (str, optional): The name of the table to process.
input_filters (dict, optional): Filters to be applied during prediction.
mlflow_transform_filters (dict, optional): Filters to be applied during prediction.
mlflow_predict_filters (dict, optional): Filters to be applied during prediction.
model_retention (int, optional): The model retention period in minutes.
path_priority (list[str]): The path priority.
Returns:
None
Raises:
Exception: If any of the required parameters are missing or if the workflow fails.
"""
await workflow.execute_local_activity_method(
Activities.prepare_activity,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
data = await workflow.execute_local_activity_method(
Activities.load_custom_query,
input_data['query'],
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
# Prepare input for prediction_process workflow
prediction_input = {
'data': data,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'model_retention': input_data.get('model_retention', 60),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {})
}
await workflow.execute_child_workflow(
'prediction_process', prediction_input)

View File

@@ -1,95 +0,0 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from datetime import timedelta
from laborious.utils.policies import retry_policy
@workflow.defn(name="format_and_export_prediction")
class FormatAndExportPrediction():
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
This workflow formats and exports predictions based on path_flag:
- If path_flag is None: formats prediction
using input data, timestamp, model_id and confidence
- If path_flag exists: creates default prediction
with timestamp, model_id, confidence and comment
Finally exports formatted prediction to postgres table
Args:
input_data(dict[str, Any]): The input data for the workflow.
Contains the following keys:
path_flag(str): The path flag to determine the type of prediction to format
data(dict[str, Any]): The data to format
prediction_confidence(float): The prediction confidence to be registered
timestamp(str): The timestamp of the prediction, synchronized with the data
model_id(int): The model id of the prediction
model_name(str): The model name of the prediction
model_retention(str): The model retention of the prediction
comment(str): The comment to be registered
schema(str): The schema of the prediction
table_name(str): The table name of the prediction
opc_output_config(dict[str, Any]): The opc output config of the prediction
Returns:
bool: True if the workflow was successful, False otherwise.
"""
path_flag = input_data['path_flag']
data = input_data['data']
prediction_confidence = input_data['prediction_confidence']
if path_flag is None:
# proceed with formatting and exporting
prediction = await workflow.execute_local_activity_method(
Activities.format_prediction,
{
'data': data,
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
else:
# create default prediction
prediction = await workflow.execute_local_activity_method(
Activities.format_default_prediction,
{
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'comment': input_data['comment']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
# write to postgres
postgres_holder = workflow.execute_activity_method(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
# write to opc
opc_holder = workflow.execute_activity_method(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': prediction
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
await postgres_holder
await opc_holder

View File

@@ -1,233 +0,0 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from laborious.utils.policies import retry_policy
from datetime import timedelta
@workflow.defn(name="prediction_process")
class PredictionProcess():
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
This workflow runs a prediction process based on the input data.
The workflow executes in two main steps:
1. Prepares the activity with schedule and model information
2. Loads data using a custom query and executes the prediction process
Args:
input_data (dict[str, Any]): The input data for the workflow.
Contains the following keys:
data (dict[str, Any]): The data to be used for the prediction.
schema (str): The schema of the table.
table_name (str): The name of the table.
model_id (int): The id of the model.
input_filters (dict, optional): Filters to be applied during prediction.
mlflow_transform_filters (dict, optional): Filters to be applied during prediction.
mlflow_predict_filters (dict, optional): Filters to be applied during prediction.
model_name (str): The name of the model.
model_retention (int, optional): The model retention period in minutes.
path_priority (list[str]): The path priority.
opc_output_config (dict[str, Any]): The opc output config of the prediction.
Returns:
None
Raises:
Exception: If any of the required parameters are missing or if the workflow fails.
"""
data = input_data['data']
model_id = input_data['model_id']
model_name = input_data['model_name']
model_retention = input_data['model_retention']
last_timestamp = await workflow.execute_local_activity_method(
Activities.get_last_timestamp,
{
'data': data
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': data,
'path_priority': input_data['path_priority']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
response_data = await workflow.execute_local_activity_method(
Activities.request_transform,
{
'data': data,
'model_name': model_name,
'model_retention': model_retention
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': response_data,
'type': 'transform',
'path_priority': input_data['path_priority']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
transformed_data = response_data['content']
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_content_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': transformed_data,
'type': 'transform',
'path_priority': input_data['path_priority']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
response_data = await workflow.execute_local_activity_method(
Activities.request_predict,
{
'data': transformed_data,
'model_name': model_name,
'model_retention': model_retention
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_predict_filters'],
'data': response_data,
'type': 'predict',
'path_priority': input_data['path_priority']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
await workflow.execute_child_workflow(
'format_and_export_prediction',
{
'path_flag': path_flag,
'data': response_data['content'],
'prediction_confidence': confidence,
'timestamp': last_timestamp,
'model_id': model_id,
'model_name': model_name,
'model_retention': model_retention,
'opc_output_config': input_data['opc_output_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'comment': comment
}
)
async def path_flag_handler(self, data: dict[str, Any], path_flag: str,
input_data: dict[str, Any], confidence: int,
last_timestamp: str, comment: str):
"""
This function handles the path flag and the confidence of the prediction.
It returns True if the prediction should be stopped. If path_flag is 'repeat',
it repeats the last prediction.
If path_flag is 'continue', it calls the write workflow. If path_flag is 'stop',
it stops the prediction process.
Args:
data (dict[str, Any]): The data to be used for the prediction.
path_flag (str): The path flag to determine the type of prediction to format
confidence (int): The confidence of the prediction
schema (str): The schema of the prediction
table_name (str): The table name of the prediction
model_id (int): The model id of the prediction
last_timestamp (str): The timestamp of the last prediction
model_name (str): The model name of the prediction
model_retention (int): The model retention of the prediction
comment (str): The comment of the prediction
Returns:
bool: True if the prediction should be stopped, False otherwise.
"""
schema = input_data['schema']
table_name = input_data['table_name']
model_id = input_data['model_id']
model_name = input_data['model_name']
model_retention = input_data['model_retention']
path_flag = path_flag.upper() if path_flag else None
if path_flag == 'STOP':
return True
elif path_flag == 'REPEAT':
# repeat last prediction
await workflow.execute_activity_method(
Activities.repeat_last_prediction,
{
'schema': schema,
'table_name': table_name,
'model_id': model_id
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
return True
elif path_flag == 'CONTINUE':
# call write workflow
await workflow.execute_child_workflow(
'format_and_export_prediction',
{
'path_flag': path_flag,
'data': data,
'prediction_confidence': confidence,
'timestamp': last_timestamp,
'model_id': model_id,
'model_name': model_name,
'model_retention': model_retention,
'schema': schema,
'table_name': table_name,
'comment': comment,
'opc_output_config': input_data['opc_output_config']
}
)
return True
return False

View File

@@ -0,0 +1,47 @@
from temporalio import activity, workflow
from temporalio.client import Client
with workflow.unsafe.imports_passed_through():
from orchestrator.activities.couchbase import Couchbase
from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.activities.slot_manager import SlotManager
from typing import Any
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
class Activities(Couchbase, TemporalManager, SlotManager):
def __init__(self,
temporal_client: Client,
couchbase_config: dict[str, Any],
redis_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler):
# Initialize parent classes
Couchbase.__init__(self, connection_string=couchbase_config['connection_string'],
username=couchbase_config['username'],
password=couchbase_config['password'],
logger=logger,
notification_handler=notification_handler)
TemporalManager.__init__(self,
temporal_client=temporal_client,
logger=logger,
notification_handler=notification_handler)
SlotManager.__init__(self,
host=redis_config['host'],
port=redis_config['port'],
username=redis_config['username'],
password=redis_config['password'],
logger=logger,
notification_handler=notification_handler)
@activity.defn(name="prepare_activity")
async def prepare_activity(self, input_data: dict[str, Any]):
await super().prepare_activity(input_data)
def shutdown(self):
Couchbase.shutdown(self)

View File

@@ -0,0 +1,94 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
from logging import Logger
from datetime import timedelta
import json
import traceback
from couchbase.auth import PasswordAuthenticator
from couchbase.cluster import Cluster
from couchbase.options import ClusterOptions
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity
class Couchbase(BaseActivity):
def __init__(self, connection_string: str, username: str,
password: str, logger: Logger,
notification_handler: NotificationHandler):
self.connection_string = connection_string
self.username = username
self.password = password
logger.info("Initializing Couchbase connection...")
self.cluster = Cluster(
connection_string,
ClusterOptions(
authenticator=PasswordAuthenticator(
username=username,
password=password
)
)
)
logger.info("Awaiting Couchbase connection...")
self.cluster.wait_until_ready(timeout=timedelta(seconds=10))
logger.info("Couchbase connection ready")
BaseActivity.__init__(self,
logger=logger,
notification_handler=notification_handler)
def shutdown(self):
try:
self.cluster.close()
except Exception as e:
self.logger.error("Failed to close Couchbase connection: %s", e)
def __del__(self):
self.shutdown()
@activity.defn(name="load_query_from_couchbase")
async def load_query_from_couchbase(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Load a query from couchbase
Args:
input_data (dict[str, Any]): The input data containing the query to execute
Returns:
list[dict[str, Any]]: The result of the query
"""
query = input_data['query']
self.logger.info("Executing couchbase query: %s", query)
try:
result = self.cluster.query(query)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id="COUCHBASE_LOAD_QUERY_ERROR",
message=f"Failed to execute couchbase query: {e}",
block="load_query_from_couchbase",
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.logger.error(trace)
raise e
rows = []
for row in result.rows():
rows.append(row)
self.logger.info("Fetched %d rows from couchbase", len(rows))
self.logger.debug("Rows: \n %s",
json.dumps(rows, indent=4, sort_keys=True))
return rows

View File

@@ -0,0 +1,130 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.base import BaseActivity
from typing import Any
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
class Formatters(BaseActivity):
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
BaseActivity.__init__(self, logger=logger,
notification_handler=notification_handler)
@activity.defn(name="process_schedules")
async def process_schedules(self, input_data: dict[str, Any]):
pipelines = input_data['pipelines']
schedule_config = {}
for pipeline in pipelines:
if pipeline['workflow_type'] == 'scouter':
schedule_config[pipeline['schedule_name']] = scouter(pipeline)
return schedule_config
def common_config(config: dict[str, Any]):
return {
"workflow_type": "scouter",
"schedule_name": config['schedule_name'],
"frequency": config.get('frequency', '1m'),
"max_retry_policy": config.get('max_retry_policy', 1),
"model_id": config['model_id'],
"model_name": config['model_name'],
}
def scouter(config: dict[str, Any]):
filters = {}
for f in config['filters']:
filters[f['filter_name']] = {
"policy": f['policy']
}
tags = {}
for tag in config['read_tags']:
tags[tag['tag_name']] = {
"aggr_func": tag.get('aggr_func', 'lts'),
"data_range": tag.get('data_range', [-100, 100])
}
return {
**common_config(config),
"topic": f"raw_{config['schedule_name']}",
"trigger_laborious": False,
"filters": filters,
"schema": "sientia_data",
"table_name": "laborious_data",
"retention_time": config.get('tag_retention_minutes', 60) * 60,
"model_tags": tags
}
def overlap_filter_config(base_filter_config: dict[str, Any], config: dict[str, Any]):
for fil in config['filters']:
base_filter_config[fil['filter_name']] = {
"policy": fil['policy'],
"config": fil.get('config', {})
}
return base_filter_config
def predictions_batch(config: dict[str, Any]):
tags = {}
for tag in config['write_tags']:
if tag['server_name'] not in tags:
tags[tag['server_name']] = {}
tag_type = tag['type']
if tag_type == 'prediction' or tag_type == 'confidence':
tag_type_str = f"{tag_type}_tags"
if tag_type_str not in tags[tag['server_name']]:
tags[tag['server_name']][tag_type_str] = {}
tags[tag['server_name']][tag_type_str][tag['addr']] = {
"data_type": tag.get('data_type', 'float'),
}
path_priority = config.get('path_priority', ["STOP", "CONTINUE", "REPEAT"])
for priority in path_priority[:]:
if priority not in ["STOP", "CONTINUE", "REPEAT"]:
path_priority.remove(priority)
if len(path_priority) != 3:
for priority in ["STOP", "CONTINUE", "REPEAT"]:
if priority not in path_priority:
path_priority.append(priority)
return {
**common_config(config),
"query": config['query'],
"schema": "sientia_data",
"table_name": "predictions",
"retention_time": config.get('model_retention_minutes', 60) * 60,
"opc_output_config": tags,
"input_filters": overlap_filter_config({
"EMPTY_DATA": {
"policy": "STOP"
}
}, config['input_filters']),
"mlflow_transform_filters": overlap_filter_config({
"API_ERROR": {
"policy": "STOP"
}
}, config['mlflow_transform_filters']),
"mlflow_predict_filters": overlap_filter_config({
"API_ERROR": {
"policy": "STOP"
}
}, config['mlflow_predict_filters']),
"path_priority": path_priority
}

View File

@@ -0,0 +1,73 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
import json
from logging import Logger
from sientia_do.temporal.activities.redis_base import Redis
from sientia_do.notifications.handlers import NotificationHandler
class SlotManager(Redis):
def __init__(self, host: str, port: int,
username: str, password: str,
logger: Logger, notification_handler: NotificationHandler):
Redis.__init__(self, host, port, username,
password, logger, notification_handler)
@activity.defn(name="load_opc_slots")
async def load_opc_slots(self) -> dict[str, Any]:
"""
Load all OPC slots from Redis
Returns:
dict[str, Any]: A dictionary of OPC slots
"""
self.logger.info("Loading OPC slots...")
opc_slots = {}
slot_keys = self.redis_client.keys("slot:opc_tags:*")
if slot_keys:
decoded_keys = [key.decode('utf-8') for key in slot_keys]
values = self.redis_client.mget(decoded_keys)
for i, key in enumerate(decoded_keys):
value = values[i]
if value is not None:
try:
opc_slots[key] = value.decode('utf-8')
except (UnicodeDecodeError, AttributeError):
opc_slots[key] = value
else:
opc_slots[key] = None
self.logger.info(f"Loaded {len(opc_slots)} OPC slots")
self.logger.debug("OPC slots: \n %s",
json.dumps(opc_slots, indent=4, sort_keys=True))
return opc_slots
@activity.defn(name="load_active_ingestors")
async def load_active_ingestors(self) -> list[str]:
"""
Load all active ingestors from Redis
Returns:
list[str]: A list of active ingestors
"""
self.logger.info("Loading active ingestors...")
active_ingestors = self.redis_client.keys("heartbeat:ingestor:*")
self.logger.info(f"Loaded {len(active_ingestors)} active ingestors")
self.logger.debug("Active ingestors: \n %s", active_ingestors)
return [ingestor.decode('utf-8') for ingestor in active_ingestors]

View File

@@ -0,0 +1,66 @@
from temporalio import activity, workflow
from temporalio.client import Client
with workflow.unsafe.imports_passed_through():
from typing import Any
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.temporal.activities.base import BaseActivity
from google.protobuf.json_format import MessageToDict
import base64
import json
class TemporalManager(BaseActivity):
def __init__(self, temporal_client: Client, logger: Logger,
notification_handler: NotificationHandler):
self.temporal_client = temporal_client
BaseActivity.__init__(self,
logger=logger,
notification_handler=notification_handler)
@activity.defn(name="load_schedule")
async def load_schedule(self) -> dict[str, Any]:
"""
Load all orchestrated schedules from Temporal. Filters by search attribute
"Orchestrated" set to "true" and returns a dictionary of schedule_id:
{frequency, data, handle}
Returns:
dict[str, Any]: A dictionary of orchestrated schedules
"""
self.logger.info("Getting orchestrated schedules...")
orchestrated_schedules = {}
async for schedule in await self.temporal_client.list_schedules():
search_attrs = getattr(schedule, "search_attributes", {})
if search_attrs.get("Orchestrated", ["false"]) == ["true"]:
schedule_id = schedule.id
handle = self.temporal_client.get_schedule(schedule_id)
desc = await handle.describe()
for arg in desc.schedule.action.args:
data = MessageToDict(arg)['data']
data = base64.b64decode(data).decode('utf-8')
frequency = desc.schedule.spec.intervals[0].every.seconds
orchestrated_schedules[schedule_id] = {
'frequency': frequency,
'data': json.loads(data),
'handle': handle
}
self.logger.info("Found %d orchestrated schedules",
len(orchestrated_schedules))
self.logger.debug("Orchestrated schedules: %s",
orchestrated_schedules)
return orchestrated_schedules

View File

@@ -0,0 +1,69 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from orchestrator.activities.activities import Activities
from typing import Any
from datetime import timedelta
from sientia_do.temporal.utils.policies import retry_policy
@workflow.defn(name="orchestrator")
class Orchestrator:
@workflow.run
async def run(self, input_data: dict[str, Any]):
input_data['workflow_name'] = 'orchestrator'
await workflow.execute_local_activity_method(
Activities.prepare_activity,
{
'workflow_name': input_data['workflow_name'],
'schedule_name': input_data['schedule_name'],
'model_name': '-',
'model_id': '-'
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
pipeline_config_handler = workflow.execute_local_activity_method(
Activities.load_query_from_couchbase,
{
'query': input_data['pipelines_query']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
opc_servers_handler = workflow.execute_local_activity_method(
Activities.load_query_from_couchbase,
{
'query': input_data['opc_servers_query']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
orchestrated_schedules_handler = workflow.execute_local_activity_method(
Activities.load_schedule,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
slot_config_handler = workflow.execute_local_activity_method(
Activities.load_opc_slots,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
active_ingestors_handler = workflow.execute_local_activity_method(
Activities.load_active_ingestors,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
pipeline_config = await pipeline_config_handler
orchestrated_schedules = await orchestrated_schedules_handler
slot_config = await slot_config_handler
opc_servers = await opc_servers_handler
active_ingestors = await active_ingestors_handler

View File

@@ -2,4 +2,5 @@ temporalio
psycopg2-binary psycopg2-binary
sqlalchemy sqlalchemy
redis redis
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git couchbase
/home/grezewave/Documents/projects/sientia/sientia-dataops-library/

35
samples.json Normal file
View File

@@ -0,0 +1,35 @@
{
"models": {
"1": {
"name": "Demo Model-Demo2"
}
},
"pipelines": {
"1": {
"name": "scouter-opcua-pipeline",
"model_id": 1,
"workflow_type": "scouter",
"frequency": "5s",
"max_retry_policy": 1,
"read_tags": [
{
"tag_name": "Counter",
"aggr_func": "avg",
"data_range": [-100, 100]
}
],
"filters": [
{
"filter_name": "OUT_OF_BOUNDS_FILTER",
"policy": "DISCARD"
},
{
"filter_name": "NULL_VALUES_FILTER",
"policy": "DISCARD"
}
],
"tag_retention_minutes": 60
}
}
}

View File

@@ -1,30 +0,0 @@
# syntax=docker/dockerfile:1.4
FROM python:3.11-slim
# Enable use of SSH agent/socket
# This line enables SSH during build
# (don't forget the syntax header above)
RUN apt-get update && apt-get install -y git openssh-client && rm -rf /var/lib/apt/lists/*
# Use build-time SSH mount for Git clone
# The SSH key will NOT remain in the image
# IMPORTANT: this block requires BuildKit
# and the --ssh flag during docker build
# SSH config to skip host key check (safe in CI/local dev)
RUN mkdir -p /root/.ssh && echo "StrictHostKeyChecking no" > /root/.ssh/config
WORKDIR /app
# Clone using SSH
ARG GIT_REPO
ARG GIT_BRANCH=main
# Mount SSH key just for this RUN
RUN --mount=type=ssh git clone --branch ${GIT_BRANCH} ${GIT_REPO} .
# Install requirements if exists
RUN if [ -f requirements.txt ]; then pip install --no-cache-dir -r requirements.txt; fi
CMD ["python", "server.py"]

876
test.ipynb Normal file
View File

@@ -0,0 +1,876 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 4,
"id": "a287fa45",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Initializing Couchbase connection...\n",
"Awaiting Couchbase connection...\n",
"Couchbase connection ready\n"
]
}
],
"source": [
"from orchestrator.activities.couchbase import Couchbase\n",
"from unittest.mock import MagicMock\n",
"logger = MagicMock(info=MagicMock(side_effect=print), debug=MagicMock(side_effect=print))\n",
"couchbase = Couchbase(\n",
" connection_string=\"couchbase://localhost\",\n",
" username=\"sientia\",\n",
" password=\"sientia\",\n",
" logger=logger,\n",
" notification_handler=MagicMock()\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "d9a0f9c4",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Executing couchbase query: %s \n",
"SELECT\n",
" pipelines.*,\n",
" models as model\n",
"FROM\n",
" `pipelines`\n",
"JOIN\n",
" `models` ON KEYS pipelines.model_id;\n",
"\n",
"Fetched %d rows from couchbase 1\n",
"Rows: \n",
" %s [\n",
" {\n",
" \"filters\": [\n",
" {\n",
" \"filter_name\": \"OUT_OF_BOUNDS_FILTER\",\n",
" \"policy\": \"DISCARD\"\n",
" },\n",
" {\n",
" \"filter_name\": \"NULL_VALUES_FILTER\",\n",
" \"policy\": \"DISCARD\"\n",
" }\n",
" ],\n",
" \"frequency\": \"5s\",\n",
" \"max_retry_policy\": 1,\n",
" \"model\": {\n",
" \"name\": \"Demo Model-Demo2\"\n",
" },\n",
" \"model_id\": \"1\",\n",
" \"name\": \"scouter-opcua-pipeline\",\n",
" \"read_tags\": [\n",
" {\n",
" \"aggr_func\": \"avg\",\n",
" \"data_range\": [\n",
" -100,\n",
" 100\n",
" ],\n",
" \"tag_name\": \"Counter\"\n",
" }\n",
" ],\n",
" \"tag_retention_minutes\": 60,\n",
" \"workflow_type\": \"scouter\"\n",
" }\n",
"]\n",
"[{'filters': [{'filter_name': 'OUT_OF_BOUNDS_FILTER', 'policy': 'DISCARD'}, {'filter_name': 'NULL_VALUES_FILTER', 'policy': 'DISCARD'}], 'frequency': '5s', 'max_retry_policy': 1, 'model': {'name': 'Demo Model-Demo2'}, 'model_id': '1', 'name': 'scouter-opcua-pipeline', 'read_tags': [{'aggr_func': 'avg', 'data_range': [-100, 100], 'tag_name': 'Counter'}], 'tag_retention_minutes': 60, 'workflow_type': 'scouter'}]\n"
]
}
],
"source": [
"query = \"\"\"\n",
"SELECT\n",
" pipelines.*,\n",
" models as model\n",
"FROM\n",
" `pipelines`\n",
"JOIN\n",
" `models` ON KEYS pipelines.model_id;\n",
"\"\"\"\n",
"if __name__ == \"__main__\":\n",
"\n",
"\n",
" result = await couchbase.load_query_from_couchbase({\n",
" \"query\": query\n",
" })\n",
"\n",
" print(result)\n",
" "
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "7d01f160",
"metadata": {},
"outputs": [],
"source": [
"from temporalio import client\n",
"from orchestrator.activities.temporal_manager import TemporalManager\n",
"import os\n",
"from unittest.mock import MagicMock\n",
"\n",
"host = \"localhost:7233\"\n",
"logger = MagicMock(info=MagicMock(side_effect=print), debug=MagicMock(side_effect=print))\n",
"\n",
"temporal_client = await client.Client.connect(\n",
" target_host=host,\n",
" namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')\n",
")\n",
"\n",
"manager = TemporalManager(temporal_client=temporal_client,\n",
" logger=logger,\n",
" notification_handler=MagicMock())\n",
" "
]
},
{
"cell_type": "code",
"execution_count": 81,
"id": "bb750ae6",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<temporalio.client.ScheduleHandle at 0x76f314e5df90>"
]
},
"execution_count": 81,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import asyncio\n",
"from datetime import timedelta\n",
"from temporalio.client import (\n",
" Client,\n",
" Schedule,\n",
" ScheduleActionStartWorkflow,\n",
" ScheduleIntervalSpec,\n",
" ScheduleSpec,\n",
")\n",
"from temporalio.common import TypedSearchAttributes, SearchAttributeKey, SearchAttributePair\n",
"\n",
"\n",
"customer_id_key = SearchAttributeKey.for_keyword(\"Orchestrated\")\n",
"search_attributes = TypedSearchAttributes([\n",
" SearchAttributePair(customer_id_key, \"true\")\n",
"])\n",
"await temporal_client.create_schedule(\n",
" \"meu-schedule-id5\",\n",
" Schedule(\n",
" action=ScheduleActionStartWorkflow(\n",
" 'scouter-test2',\n",
" {\n",
" 'args': {\n",
" 'arg1': 'value1'\n",
" }\n",
" },\n",
" id=\"workflow-id-unico\",\n",
" task_queue=\"nome-da-task-queue\",\n",
" ),\n",
" spec=ScheduleSpec(\n",
" intervals=[ScheduleIntervalSpec(every=timedelta(minutes=10))]\n",
" )\n",
" ),\n",
" search_attributes=search_attributes,\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 86,
"id": "1bd82225",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Getting orchestrated schedules...\n",
"Schedule: %s ScheduleListDescription(id='meu-schedule-id5', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test2'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=<SearchAttributeIndexedValueType.TEXT: 1>, _value_type=<class 'str'>), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x76f32ca5cf50>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x76f32ca5cf90>), raw_entry=schedule_id: \"meu-schedule-id5\"\n",
"search_attributes {\n",
" indexed_fields {\n",
" key: \"Orchestrated\"\n",
" value {\n",
" metadata {\n",
" key: \"type\"\n",
" value: \"Text\"\n",
" }\n",
" metadata {\n",
" key: \"encoding\"\n",
" value: \"json/plain\"\n",
" }\n",
" data: \"\\\"true\\\"\"\n",
" }\n",
" }\n",
"}\n",
"info {\n",
" spec {\n",
" interval {\n",
" interval {\n",
" seconds: 600\n",
" }\n",
" }\n",
" }\n",
" workflow_type {\n",
" name: \"scouter-test2\"\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548800\n",
" }\n",
" future_action_times {\n",
" seconds: 1748549400\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550000\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550600\n",
" }\n",
" future_action_times {\n",
" seconds: 1748551200\n",
" }\n",
"}\n",
")\n",
"Search attributes: %s {'Orchestrated': ['true']}\n",
"Schedule: %s ScheduleListDescription(id='meu-schedule-id4', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test2'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=<SearchAttributeIndexedValueType.TEXT: 1>, _value_type=<class 'str'>), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x76f32ca5cf50>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x76f32ca5cf90>), raw_entry=schedule_id: \"meu-schedule-id4\"\n",
"search_attributes {\n",
" indexed_fields {\n",
" key: \"Orchestrated\"\n",
" value {\n",
" metadata {\n",
" key: \"type\"\n",
" value: \"Text\"\n",
" }\n",
" metadata {\n",
" key: \"encoding\"\n",
" value: \"json/plain\"\n",
" }\n",
" data: \"\\\"true\\\"\"\n",
" }\n",
" }\n",
"}\n",
"info {\n",
" spec {\n",
" interval {\n",
" interval {\n",
" seconds: 600\n",
" }\n",
" }\n",
" }\n",
" workflow_type {\n",
" name: \"scouter-test2\"\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548800\n",
" }\n",
" future_action_times {\n",
" seconds: 1748549400\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550000\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550600\n",
" }\n",
" future_action_times {\n",
" seconds: 1748551200\n",
" }\n",
"}\n",
")\n",
"Search attributes: %s {'Orchestrated': ['true']}\n",
"Schedule: %s ScheduleListDescription(id='meu-schedule-id3', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test2'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=<SearchAttributeIndexedValueType.TEXT: 1>, _value_type=<class 'str'>), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x76f32ca5cf50>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x76f32ca5cf90>), raw_entry=schedule_id: \"meu-schedule-id3\"\n",
"search_attributes {\n",
" indexed_fields {\n",
" key: \"Orchestrated\"\n",
" value {\n",
" metadata {\n",
" key: \"type\"\n",
" value: \"Text\"\n",
" }\n",
" metadata {\n",
" key: \"encoding\"\n",
" value: \"json/plain\"\n",
" }\n",
" data: \"\\\"true\\\"\"\n",
" }\n",
" }\n",
"}\n",
"info {\n",
" spec {\n",
" interval {\n",
" interval {\n",
" seconds: 600\n",
" }\n",
" }\n",
" }\n",
" workflow_type {\n",
" name: \"scouter-test2\"\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548800\n",
" }\n",
" future_action_times {\n",
" seconds: 1748549400\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550000\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550600\n",
" }\n",
" future_action_times {\n",
" seconds: 1748551200\n",
" }\n",
"}\n",
")\n",
"Search attributes: %s {'Orchestrated': ['true']}\n",
"Schedule: %s ScheduleListDescription(id='meu-schedule-id2', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test2'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 50, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 50, 0, 37623, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='workflow-id-unico-2025-05-29T19:50:00Z', first_execution_run_id='01971d98-265a-7285-b46f-cf09bcf2d301'))], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=<SearchAttributeIndexedValueType.TEXT: 1>, _value_type=<class 'str'>), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x76f32ca5cf50>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x76f32ca5cf90>), raw_entry=schedule_id: \"meu-schedule-id2\"\n",
"search_attributes {\n",
" indexed_fields {\n",
" key: \"Orchestrated\"\n",
" value {\n",
" metadata {\n",
" key: \"type\"\n",
" value: \"Text\"\n",
" }\n",
" metadata {\n",
" key: \"encoding\"\n",
" value: \"json/plain\"\n",
" }\n",
" data: \"\\\"true\\\"\"\n",
" }\n",
" }\n",
"}\n",
"info {\n",
" spec {\n",
" interval {\n",
" interval {\n",
" seconds: 600\n",
" }\n",
" }\n",
" }\n",
" workflow_type {\n",
" name: \"scouter-test2\"\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548200\n",
" }\n",
" actual_time {\n",
" seconds: 1748548200\n",
" nanos: 37623285\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"workflow-id-unico-2025-05-29T19:50:00Z\"\n",
" run_id: \"01971d98-265a-7285-b46f-cf09bcf2d301\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548800\n",
" }\n",
" future_action_times {\n",
" seconds: 1748549400\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550000\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550600\n",
" }\n",
" future_action_times {\n",
" seconds: 1748551200\n",
" }\n",
"}\n",
")\n",
"Search attributes: %s {'Orchestrated': ['true']}\n",
"Schedule: %s ScheduleListDescription(id='meu-schedule-id', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter-test'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 0, 0, 37788, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='workflow-id-unico-2025-05-29T19:00:00Z', first_execution_run_id='01971d6a-5fa1-70f8-8371-60ad11587b77'))], next_action_times=[datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 10, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 20, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 40, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[SearchAttributePair(key=_SearchAttributeKey(_name='Orchestrated', _indexed_value_type=<SearchAttributeIndexedValueType.TEXT: 1>, _value_type=<class 'str'>), value='true')]), search_attributes={'Orchestrated': ['true']}, data_converter=DataConverter(payload_converter_class=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x76f32ca5cf50>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x76f32ca5cf90>), raw_entry=schedule_id: \"meu-schedule-id\"\n",
"search_attributes {\n",
" indexed_fields {\n",
" key: \"Orchestrated\"\n",
" value {\n",
" metadata {\n",
" key: \"type\"\n",
" value: \"Text\"\n",
" }\n",
" metadata {\n",
" key: \"encoding\"\n",
" value: \"json/plain\"\n",
" }\n",
" data: \"\\\"true\\\"\"\n",
" }\n",
" }\n",
"}\n",
"info {\n",
" spec {\n",
" interval {\n",
" interval {\n",
" seconds: 600\n",
" }\n",
" }\n",
" }\n",
" workflow_type {\n",
" name: \"scouter-test\"\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748545200\n",
" }\n",
" actual_time {\n",
" seconds: 1748545200\n",
" nanos: 37788631\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"workflow-id-unico-2025-05-29T19:00:00Z\"\n",
" run_id: \"01971d6a-5fa1-70f8-8371-60ad11587b77\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548800\n",
" }\n",
" future_action_times {\n",
" seconds: 1748549400\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550000\n",
" }\n",
" future_action_times {\n",
" seconds: 1748550600\n",
" }\n",
" future_action_times {\n",
" seconds: 1748551200\n",
" }\n",
"}\n",
")\n",
"Search attributes: %s {'Orchestrated': ['true']}\n",
"Schedule: %s ScheduleListDescription(id='laborious_test', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='predictions_batch'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=60), offset=datetime.timedelta(0))], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 53, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 53, 0, 35750, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:53:00Z', first_execution_run_id='01971d9a-e57f-700b-953b-bdb3b05810e2')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 54, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 54, 0, 35780, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:54:00Z', first_execution_run_id='01971d9b-cfde-7f39-8b38-b7e62fee1f82')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 55, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 55, 0, 34329, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:55:00Z', first_execution_run_id='01971d9c-ba3c-7d27-9db7-72759ebabc76')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 56, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 56, 0, 49214, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:56:00Z', first_execution_run_id='01971d9d-a4a9-7b55-a77b-fd450e768ccf')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 57, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 57, 0, 36109, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='laborious_test-2025-05-29T19:57:00Z', first_execution_run_id='01971d9e-8eff-7608-9b10-0ed1875df9fe'))], next_action_times=[datetime.datetime(2025, 5, 29, 19, 58, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 59, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 1, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 20, 2, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[]), search_attributes={}, data_converter=DataConverter(payload_converter_class=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x76f32ca5cf50>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x76f32ca5cf90>), raw_entry=schedule_id: \"laborious_test\"\n",
"info {\n",
" spec {\n",
" interval {\n",
" interval {\n",
" seconds: 60\n",
" }\n",
" phase {\n",
" }\n",
" }\n",
" }\n",
" workflow_type {\n",
" name: \"predictions_batch\"\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548380\n",
" }\n",
" actual_time {\n",
" seconds: 1748548380\n",
" nanos: 35750962\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"laborious_test-2025-05-29T19:53:00Z\"\n",
" run_id: \"01971d9a-e57f-700b-953b-bdb3b05810e2\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548440\n",
" }\n",
" actual_time {\n",
" seconds: 1748548440\n",
" nanos: 35780414\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"laborious_test-2025-05-29T19:54:00Z\"\n",
" run_id: \"01971d9b-cfde-7f39-8b38-b7e62fee1f82\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548500\n",
" }\n",
" actual_time {\n",
" seconds: 1748548500\n",
" nanos: 34329717\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"laborious_test-2025-05-29T19:55:00Z\"\n",
" run_id: \"01971d9c-ba3c-7d27-9db7-72759ebabc76\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548560\n",
" }\n",
" actual_time {\n",
" seconds: 1748548560\n",
" nanos: 49214684\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"laborious_test-2025-05-29T19:56:00Z\"\n",
" run_id: \"01971d9d-a4a9-7b55-a77b-fd450e768ccf\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548620\n",
" }\n",
" actual_time {\n",
" seconds: 1748548620\n",
" nanos: 36109875\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"laborious_test-2025-05-29T19:57:00Z\"\n",
" run_id: \"01971d9e-8eff-7608-9b10-0ed1875df9fe\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548680\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548740\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548800\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548860\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548920\n",
" }\n",
"}\n",
")\n",
"Search attributes: %s {}\n",
"Schedule: %s ScheduleListDescription(id='scouter-opcua-pipeline', schedule=ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='scouter'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=30), offset=datetime.timedelta(0))], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), info=ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 55, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 55, 0, 26130, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:55:00Z', first_execution_run_id='01971d9c-ba35-7b2a-b596-effe9939c7c7')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 55, 30, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 55, 30, 26550, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:55:30Z', first_execution_run_id='01971d9d-2f66-7012-bd77-e5809b8c4c7a')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 56, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 56, 0, 39848, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:56:00Z', first_execution_run_id='01971d9d-a4a1-7ba9-9205-21f157fe6e54')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 56, 30, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 56, 30, 39791, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:56:30Z', first_execution_run_id='01971d9e-19d2-7358-96d3-261e203faaa7')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 5, 29, 19, 57, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 5, 29, 19, 57, 0, 28221, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='scouter-opcua-pipeline-2025-05-29T19:57:00Z', first_execution_run_id='01971d9e-8ef8-72af-a903-1391b5e06c2f'))], next_action_times=[datetime.datetime(2025, 5, 29, 19, 57, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 58, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 58, 30, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 59, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 5, 29, 19, 59, 30, tzinfo=datetime.timezone.utc)]), typed_search_attributes=TypedSearchAttributes(search_attributes=[]), search_attributes={}, data_converter=DataConverter(payload_converter_class=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x76f32ca5cf50>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x76f32ca5cf90>), raw_entry=schedule_id: \"scouter-opcua-pipeline\"\n",
"info {\n",
" spec {\n",
" interval {\n",
" interval {\n",
" seconds: 30\n",
" }\n",
" phase {\n",
" }\n",
" }\n",
" }\n",
" workflow_type {\n",
" name: \"scouter\"\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548500\n",
" }\n",
" actual_time {\n",
" seconds: 1748548500\n",
" nanos: 26130839\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:55:00Z\"\n",
" run_id: \"01971d9c-ba35-7b2a-b596-effe9939c7c7\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548530\n",
" }\n",
" actual_time {\n",
" seconds: 1748548530\n",
" nanos: 26550164\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:55:30Z\"\n",
" run_id: \"01971d9d-2f66-7012-bd77-e5809b8c4c7a\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548560\n",
" }\n",
" actual_time {\n",
" seconds: 1748548560\n",
" nanos: 39848794\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:56:00Z\"\n",
" run_id: \"01971d9d-a4a1-7ba9-9205-21f157fe6e54\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548590\n",
" }\n",
" actual_time {\n",
" seconds: 1748548590\n",
" nanos: 39791709\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:56:30Z\"\n",
" run_id: \"01971d9e-19d2-7358-96d3-261e203faaa7\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n",
" }\n",
" recent_actions {\n",
" schedule_time {\n",
" seconds: 1748548620\n",
" }\n",
" actual_time {\n",
" seconds: 1748548620\n",
" nanos: 28221068\n",
" }\n",
" start_workflow_result {\n",
" workflow_id: \"scouter-opcua-pipeline-2025-05-29T19:57:00Z\"\n",
" run_id: \"01971d9e-8ef8-72af-a903-1391b5e06c2f\"\n",
" }\n",
" start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548650\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548680\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548710\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548740\n",
" }\n",
" future_action_times {\n",
" seconds: 1748548770\n",
" }\n",
"}\n",
")\n",
"Search attributes: %s {}\n",
"Found %d orchestrated schedules 5\n"
]
}
],
"source": [
"schedules = await manager.load_schedule({})"
]
},
{
"cell_type": "code",
"execution_count": 87,
"id": "a4c777dd",
"metadata": {},
"outputs": [],
"source": [
"from google.protobuf.json_format import MessageToDict\n",
"import base64\n",
"import json\n",
"\n",
"\n",
"schedules_config = {}\n",
"for schedule in schedules:\n",
" id = schedule.id\n",
"\n",
" handle = temporal_client.get_schedule_handle(id)\n",
"\n",
" desc = await handle.describe()\n",
"\n",
" for arg in desc.schedule.action.args:\n",
" data = MessageToDict(arg)['data']\n",
" data = base64.b64decode(data).decode('utf-8')\n",
"\n",
" frequency = desc.schedule.spec.intervals[0].every.seconds\n",
"\n",
" schedules_config[id] = {\n",
" 'frequency': frequency,\n",
" 'data': json.loads(data),\n",
" 'handle': handle\n",
" }\n",
" \n",
" \n",
" \n",
" "
]
},
{
"cell_type": "code",
"execution_count": 88,
"id": "988d1718",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'meu-schedule-id5': {'frequency': 600,\n",
" 'data': {'args': {'arg1': 'value1'}},\n",
" 'handle': <temporalio.client.ScheduleHandle at 0x76f30f5828d0>},\n",
" 'meu-schedule-id4': {'frequency': 600,\n",
" 'data': {},\n",
" 'handle': <temporalio.client.ScheduleHandle at 0x76f30f5cf4d0>},\n",
" 'meu-schedule-id3': {'frequency': 600,\n",
" 'data': {},\n",
" 'handle': <temporalio.client.ScheduleHandle at 0x76f3141f13d0>},\n",
" 'meu-schedule-id2': {'frequency': 600,\n",
" 'data': {},\n",
" 'handle': <temporalio.client.ScheduleHandle at 0x76f30f64c710>},\n",
" 'meu-schedule-id': {'frequency': 600,\n",
" 'data': {},\n",
" 'handle': <temporalio.client.ScheduleHandle at 0x76f31417c350>}}"
]
},
"execution_count": 88,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"schedules_config"
]
},
{
"cell_type": "code",
"execution_count": 85,
"id": "c12f5e75",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"meu-schedule-id5 {'frequency': 600, 'data': {'args': {'arg1': 'value1'}}, 'handle': <temporalio.client.ScheduleHandle object at 0x76f30f5ce590>}\n",
"meu-schedule-id4 {'frequency': 1200, 'data': {'args': {'arg1': 'value1'}}, 'handle': <temporalio.client.ScheduleHandle object at 0x76f30f723910>}\n",
"meu-schedule-id4\n",
"meu-schedule-id3 {'frequency': 600, 'data': {}, 'handle': <temporalio.client.ScheduleHandle object at 0x76f30f64fe50>}\n",
"meu-schedule-id2 {'frequency': 600, 'data': {}, 'handle': <temporalio.client.ScheduleHandle object at 0x76f31417b9d0>}\n",
"meu-schedule-id {'frequency': 600, 'data': {}, 'handle': <temporalio.client.ScheduleHandle object at 0x76f30f64dad0>}\n"
]
}
],
"source": [
"from temporalio.client import Client, ScheduleUpdateInput, ScheduleUpdate, ScheduleSpec\n",
"\n",
"\n",
"for id, schedule in schedules_config.items():\n",
" print(id, schedule)\n",
" \n",
" if schedule['frequency'] != 600:\n",
" print(id)\n",
"\n",
" handle = schedule['handle']\n",
" \n",
" async def update_schedule(input: ScheduleUpdateInput) -> ScheduleUpdate:\n",
" schedule_action = input.description.schedule.action\n",
" \n",
" if hasattr(schedule_action, 'args'):\n",
" schedule_action.args = [{}]\n",
" \n",
" # Atualiza o intervalo de execução\n",
" input.description.schedule.spec.intervals = [\n",
" ScheduleIntervalSpec(every=timedelta(minutes=10))\n",
" ]\n",
" \n",
" return ScheduleUpdate(schedule=input.description.schedule)\n",
"\n",
" await handle.update(update_schedule)\n",
" \n"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "87eb10c8",
"metadata": {},
"outputs": [],
"source": [
"from redis import Redis\n",
"\n",
"redis = Redis(host='localhost', port=6379)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "fe4ad9dc",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Chave: slot:opc_tags:2, Valor: {\n",
"\"slot2\": \"value\"\n",
"}\n",
"Chave: slot:opc_tags:1, Valor: {\n",
"\"slot1\": \"value\"\n",
"}\n"
]
}
],
"source": [
"matching_keys = redis.keys(\"slot:opc_tags:*\")\n",
"\n",
"if matching_keys:\n",
" decoded_keys = [key.decode('utf-8') for key in matching_keys]\n",
" values = redis.mget(decoded_keys)\n",
"\n",
" items = {}\n",
" for i, key in enumerate(decoded_keys):\n",
" value = values[i]\n",
" if value is not None:\n",
" try:\n",
" items[key] = value.decode('utf-8')\n",
" except (UnicodeDecodeError, AttributeError):\n",
" items[key] = value\n",
" else:\n",
" items[key] = None\n",
"\n",
" for key, value in items.items():\n",
" print(f\"Chave: {key}, Valor: {value}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "87a9dc89",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

View File

@@ -1,193 +0,0 @@
from pytest import mark
from unittest.mock import patch, MagicMock, ANY
from laborious.activities.activities import Activities
from laborious.activities.postgres import Postgres
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.opc import OPC
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
)
assert isinstance(activities, Activities)
assert isinstance(activities, Postgres)
assert isinstance(activities, MLFlow)
assert isinstance(activities, OPC)
assert isinstance(activities, Gates)
mock_postgres_init.assert_called_once_with(
ANY,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler
)
mock_mlflow_init.assert_called_once_with(
ANY,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler
)
mock_opc_init.assert_called_once_with(
ANY,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler
)
mock_gates_init.assert_called_once_with(
ANY,
logger=logger,
notification_handler=notification_handler
)
@mark.asyncio
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
async def test_prepare_activity(_mock_opc_init,
_mock_mlflow_init, _mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
)
input_data = {
'workflow_name': 'test-workflow-name',
'schedule_name': 'test-schedule-name',
'model_name': 'test-model-name',
'model_id': 'test-model-id'
}
await activities.prepare_activity(input_data)
assert activities.notification_handler.base_notification.pipeline_name == input_data[
'workflow_name']
assert activities.notification_handler.base_notification.schedule_name == input_data[
'schedule_name']
assert activities.notification_handler.base_notification.model_name == input_data[
'model_name']
assert activities.notification_handler.base_notification.model_id == input_data[
'model_id']
@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
@patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
@patch('laborious.activities.activities.OPC', return_value=MagicMock())
def test_shutdown(mock_opc_init,
_mock_mlflow_init, mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
)
activities.shutdown()
mock_opc_init.shutdown.assert_called_once()
mock_postgres_init.close.assert_called_once()

View File

@@ -1,35 +0,0 @@
from unittest.mock import MagicMock
from laborious.activities.base import BaseActivity
from pytest import fixture, mark
from sientia_do.notifications.models import Notification
@fixture
def base_activity():
return BaseActivity(
logger=MagicMock(),
notification_handler=MagicMock(),
)
@mark.asyncio
async def test_prepare_activity(base_activity):
base_activity.notification_handler.base_notification = Notification(
project="project",
pipeline="pipeline",
trigger="-",
model_name="-",
model_id="-",
)
await base_activity.prepare_activity({
'workflow_name': 'test_workflow',
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id'
})
assert base_activity.notification_handler.base_notification.schedule_name == "test_schedule"
assert base_activity.notification_handler.base_notification.model_name == "test_model"
assert base_activity.notification_handler.base_notification.model_id == "test_model_id"
assert base_activity.notification_handler.base_notification.pipeline_name == "test_workflow"

View File

@@ -1,369 +0,0 @@
from unittest.mock import MagicMock, ANY, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.gates import Gates
@fixture
def gates_activity():
return Gates(
logger=MagicMock(),
notification_handler=MagicMock(),
)
@mark.asyncio
async def test_input_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'data': {'value': [1, 2, 3]},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.error.assert_called_once_with(
"Filter INVALID_FILTER not found"
)
@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
# Arrange
mock_input_filter_functions.__contains__.return_value = True
mock_input_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
'filters': {
'EMPTY_DATA': {'POLICY': 'STOP'}
},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
message="Error in filter EMPTY_DATA:{'POLICY': 'STOP'}: \n Test error",
block="input_gate",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
@mark.asyncio
async def test_input_gate_no_filters(gates_activity):
# Arrange
input_data = {
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_input_gate_with_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'EMPTY_DATA': {'POLICY': 'STOP'}
},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, "Input data with bad quality")
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
@mark.asyncio
@patch('laborious.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions,
gates_activity):
# Arrange
mock_mlflow_response_filter_functions.__contains__.return_value = True
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
@mark.asyncio
async def test_mlflow_response_gate_no_filters(gates_activity):
# Arrange
input_data = {
'filters': {},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_with_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'API_ERROR': {'POLICY': 'STOP'}
},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == ('STOP', -1, "API error occurred")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
@mark.asyncio
@patch('laborious.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions,
gates_activity):
# Arrange
mock_mlflow_content_filter_functions.__contains__.return_value = True
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
'filters': {
'API_ERROR': {'POLICY': 'STOP'}
},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
@mark.asyncio
async def test_mlflow_content_gate_no_filters(gates_activity):
# Arrange
input_data = {
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_with_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'NAN_VALUES': {'POLICY': 'STOP'}
},
'data': {'value': [None, None, None]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (
'STOP', -1, "Transformed data not passed the content filter")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called()
@mark.asyncio
async def test_format_prediction(gates_activity):
# Arrange
input_data = {
'data': {'prediction': [1], 'response_time': [0.1]},
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.9
}
# Act
result = await gates_activity.format_prediction(input_data)
# Assert
assert result['prediction'] == {0: 1}
assert result['response_time'] == {0: ANY}
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9}
assert result['prediction_status'] == {0: 'Good'}
assert result['comments'] == {0: ""}
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_format_default_prediction(gates_activity):
# Arrange
input_data = {
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.1,
'comment': 'Test comment'
}
# Act
result = await gates_activity.format_default_prediction(input_data)
# Assert
assert result['prediction'] == {0: 0}
assert result['response_time'] == {0: 0}
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.1}
assert result['prediction_status'] == {0: 'Bad'}
assert result['comments'] == {0: 'Test comment'}
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_get_last_timestamp_with_data(gates_activity):
# Arrange
input_data = {
'data': {
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
}
}
# Act
result = await gates_activity.get_last_timestamp(input_data)
# Assert
assert result == '2023-05-26 11:12:28'
@mark.asyncio
async def test_get_last_timestamp_no_data(gates_activity):
# Arrange
input_data = {
'data': {}
}
# Act
result = await gates_activity.get_last_timestamp(input_data)
# Assert
assert isinstance(result, str) # Should be a timestamp string
assert len(result) > 0

View File

@@ -1,120 +0,0 @@
from unittest.mock import MagicMock, patch
import numpy as np
from pytest import fixture, mark
from laborious.activities.mlflow import MLFlow
@patch("laborious.activities.mlflow.MLFlowRepository")
def test___init__(mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host="http://localhost",
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
logger=MagicMock(),
notification_handler=MagicMock()
)
assert mlflow.mlflow_host == "http://localhost"
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == "admin"
assert mlflow.mlflow_password == "admin"
mock_mlflow_repository.assert_called_once_with(
"http://localhost:5000", "admin", "admin"
)
@fixture
@patch("laborious.activities.mlflow.MLFlowRepository")
def mlflow(mock_mlflow_repository):
return MLFlow(
mlflow_host="http://localhost:5000",
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
logger=MagicMock(),
notification_handler=MagicMock()
)
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.max")
async def test_request_transform(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
'data': [
{'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
{'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
{'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
{'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
],
'model_name': 'test_model',
'model_retention': 30
}
# Mock the transform response
expected_response = {'prediction': [0.5, 0.6]}
mlflow.model_monitoring_repository.transform.return_value = expected_response
# Call the method
response_data = await mlflow.request_transform(input_data)
# Verify the data was correctly transformed
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.pivot.assert_called_once_with(
index='timestamp', columns='variable', values='value'
)
mock_dataframe = mock_dataframe.return_value.pivot.return_value
mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
mock_dataframe.reset_index.assert_called_once()
mock_dataframe.columns.name = None
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.transform.assert_called_once_with(
'test_model', mock_dataframe, 30
)
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.max")
async def test_request_predict(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
'data': [
{'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
{'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
{'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
{'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
],
'model_name': 'test_model',
'model_retention': 30
}
# Mock the predict response
expected_response = {'prediction': [0.5, 0.6]}
mlflow.model_monitoring_repository.predict.return_value = expected_response
# Call the method
response_data = await mlflow.request_predict(input_data)
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.replace.assert_called_once_with(
np.nan, None, inplace=True
)
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.predict.assert_called_once_with(
'test_model', mock_dataframe.return_value, 30
)

View File

@@ -1,196 +0,0 @@
from unittest.mock import patch, MagicMock, ANY, call
from pytest import fixture, mark
from laborious.activities.opc import NotificationLevel
from laborious.activities.opc import OPC
@patch("laborious.activities.opc.OpcRepository")
def test___init__(mock_opc_repository):
mock_logger = MagicMock()
server1 = MagicMock()
server2 = MagicMock()
mock_opc_repository.side_effect = [server1, server2]
mock_notification_handler = MagicMock()
servers = {
'server1': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
},
'server2': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
}
opc = OPC(
opc_servers=servers,
logger=mock_logger,
notification_handler=mock_notification_handler
)
assert opc.opc_servers == servers
assert opc.logger == mock_logger
assert opc.notification_handler == mock_notification_handler
assert opc.opc_repository['server1'] == server1
assert opc.opc_repository['server2'] == server2
mock_opc_repository.assert_has_calls([
call(
name="server1",
url="http://localhost:8080",
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
notification_handler=mock_notification_handler,
reconnection_interval=60,
),
])
mock_opc_repository.assert_has_calls([
call(
name="server2",
url="http://localhost:8080",
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
notification_handler=mock_notification_handler,
reconnection_interval=60,
)
])
server1.connect.assert_called_once()
server2.connect.assert_called_once()
@fixture
@patch("laborious.activities.opc.OpcRepository")
def opc(_mock_opc_repository):
servers = {
'server1': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
}
return OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
)
WRITE_DATA_CASES = [
('tag1', 'int', 50),
('tag2', 'float', 50.5),
('tag3', 'bool', True),
('tag4', 'string', 'test'),
]
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
def test_write_data_success(opc, tag, data_type, data):
opc.write_data(server='server1', tag=tag, data=data,
data_type=data_type, tag_type='prediction')
opc.opc_repository['server1'].write_data.assert_called_once_with(
tag, data, data_type)
def test_write_data_exception(opc):
opc.opc_repository['server1'].write_data.side_effect = Exception(
"Test error")
opc.write_data(server='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction')
opc.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="WRITE_OPC_PREDICTION_ERROR",
message="Error writing data to OPC server: Test error",
block="write_opc_data",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
opc.logger.error.assert_called_once()
@mark.asyncio
async def test_write_opc_data_success(opc):
# Arrange
input_data = {
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'opc_output_config': {
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
}
}
}
# Act
opc.write_data = MagicMock()
await opc.write_opc_data(input_data)
# Assert
opc.write_data.assert_has_calls([
call(
server='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction'
)])
opc.write_data.assert_has_calls([
call(
server='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence'
)
])
assert opc.write_data.call_count == 2
@mark.asyncio
async def test_write_opc_data_empty_config(opc):
# Arrange
input_data = {
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'opc_servers': ['server1'],
'opc_output_config': {
'prediction_tags': {},
'confidence_tags': {}
}
}
# Act
await opc.write_opc_data(input_data)
# Assert
opc.opc_repository['server1'].write_data.assert_not_called()
def test_shutdown(opc):
opc.shutdown()
opc.opc_repository['server1'].disconnect.assert_called_once()

View File

@@ -1,159 +0,0 @@
from unittest.mock import MagicMock, patch
from pytest import fixture, mark
import pandas as pd
from laborious.activities.postgres import Postgres
@fixture
@patch("laborious.activities.postgres.create_engine")
def postgres_activity(_mock_create_engine):
return Postgres(
host="localhost",
port=5432,
user="test_user",
password="test_password",
dbname="test_db",
min_connections=1,
max_connections=5,
logger=MagicMock(),
notification_handler=MagicMock()
)
@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_none_data(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_read_sql_query.return_value = None
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 0
@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_date_converted(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
mock_data['date'] = pd.to_datetime('2022-01-01')
mock_read_sql_query.return_value = mock_data
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 3
assert "column1" in result
assert "column2" in result
assert "date" in result
assert result['date'] == {0: '2022-01-01 00:00:00'}
@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_success(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
mock_read_sql_query.return_value = mock_data
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 2
assert "column1" in result
assert "column2" in result
postgres_activity.logger.info.assert_called()
@mark.asyncio
async def test_load_custom_query_error(postgres_activity):
query = "SELECT * FROM non_existent_table"
error_msg = "Table not found"
with patch("laborious.activities.postgres.read_sql_query", side_effect=ValueError(error_msg)):
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 0
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
async def test_repeat_last_prediction_success(postgres_activity):
query_items = {
"schema": "public",
"table_name": "predictions",
"model": 1
}
with patch("sqlalchemy.orm.session.Session.execute") as mock_execute:
await postgres_activity.repeat_last_prediction(query_items)
mock_execute.assert_called_once()
postgres_activity.logger.info.assert_called()
@mark.asyncio
async def test_repeat_last_prediction_error(postgres_activity):
query_items = {
"schema": "public",
"table_name": "predictions",
"model": 1
}
error_msg = "Database error"
with patch("sqlalchemy.orm.session.Session.execute", side_effect=ValueError(error_msg)):
await postgres_activity.repeat_last_prediction(query_items)
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
async def test_export_data_to_postgres_success(postgres_activity):
input_data = {
"schema": "public",
"table_name": "test_table",
"data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
}
with patch("laborious.activities.postgres.DataFrame.to_sql") as mock_to_sql:
await postgres_activity.export_data_to_postgres(input_data)
mock_to_sql.assert_called_once()
postgres_activity.logger.debug.assert_called()
@mark.asyncio
async def test_export_data_to_postgres_error(postgres_activity):
input_data = {
"schema": "public",
"table_name": "test_table",
"data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
}
error_msg = "Export failed"
with patch("laborious.activities.postgres.DataFrame.to_sql", side_effect=ValueError(error_msg)):
await postgres_activity.export_data_to_postgres(input_data)
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
async def test_close(postgres_activity):
postgres_activity.close()
postgres_activity.engine.dispose.assert_called_once()
@mark.asyncio
async def test_del(postgres_activity):
postgres_activity.close = MagicMock()
postgres_activity.__del__()
postgres_activity.close.assert_called_once()

View File

@@ -1,30 +0,0 @@
from pandas import DataFrame
from laborious.utils.filters.conditional_filters import (
filter_specific_variables_null_values,
filter_empty_data
)
def test_filter_specific_variables_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'VARIABLES': ['variable2']}) is False
def test_filter_specific_variables_null_values_with_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'VARIABLES': ['variable2']}) is True
def test_filter_empty_data():
assert filter_empty_data(DataFrame(), {}) is True
def test_filter_empty_data_with_data():
assert filter_empty_data(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
{}) is False

View File

@@ -1,22 +0,0 @@
from pandas import DataFrame
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
def test_api_error_filter_invalid_response():
assert api_error_filter(None, {}) == True # NOSONAR
def test_api_error_filter_valid_response_fail():
assert api_error_filter({'success': False}, {}) == True
def test_api_error_filter_valid_response_success():
assert api_error_filter({'success': True}, {}) == False
def test_nan_values_filter_all_nan_values():
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True
def test_nan_values_filter_no_nan_values():
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False

View File

@@ -1,278 +0,0 @@
from unittest.mock import ANY, MagicMock, patch
import numpy as np
from pandas import DataFrame
import pytest
from laborious.utils.repository.model_repository import MLFlowRepository
@pytest.fixture
def mlflow_repository():
with patch('laborious.utils.repository.model_repository.ModelServing', autospec=True) as MockModelServing:
mock_instance = MockModelServing.return_value
mock_instance.get_transformed_data = MagicMock()
repo = MLFlowRepository(
host='http://localhost:5000',
username='admin',
password='admin'
)
return repo
def test_get_current_data_df(mlflow_repository):
current_data = {
'prediction': [1, 3],
'target': [1, 1],
}
mlflow_repository.model_serving.get_transformed_data.return_value = {
'var1': [1, 2],
'var2': [2, np.nan],
}
expected = DataFrame({
'var1': [1],
'var2': [2],
'prediction': [1],
'target': [1],
})
output = mlflow_repository.get_current_data_df(current_data,
'model', 'target')
mlflow_repository.model_serving.get_transformed_data.assert_called_once_with(
'model', current_data, by='model')
diff = output.compare(expected)
assert diff.empty
def test_get_artifact(mlflow_repository):
mlflow_repository.get_artifact(
'destination', 'search_by', 'run_id', 'model', 'artifact'
)
mlflow_repository.model_serving.get_artifact.assert_called_once_with(
destination='destination',
search_by='search_by',
run_id='run_id',
model_name='model',
artifact_name='artifact'
)
def test_calculate_model_metrics(mlflow_repository):
mlflow_repository.model_serving.get_model_metrics.return_value = 'data'
real_data = 'real_data'
predictions = 'predictions'
flag = 'flag'
output = mlflow_repository.calculate_model_metrics(
real_data, predictions, flag
)
mlflow_repository.model_serving.get_model_metrics.assert_called_once_with(
reference_data=None,
real_data=real_data,
predictions=predictions,
type_flag=flag
)
assert output == 'data'
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_by_run_id(mlflow, mlflow_repository):
mlflow.get_run.return_value = MagicMock(
info=MagicMock(
experiment_id='0',
)
)
mlflow.get_experiment.return_value = MagicMock()
mlflow.get_experiment.return_value.name = 'test'
output = mlflow_repository.get_experiment_by_run_id('0')
assert output == 'test'
mlflow.get_run.assert_called_once_with('0')
mlflow.get_experiment.assert_called_once_with('0')
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_next_run_name(mlflow, mlflow_repository):
mlflow.search_runs.return_value = [1, 2, 3]
output = mlflow_repository.get_next_run_name('run')
assert output == 'run-4'
mlflow.search_runs.assert_called_once_with(
experiment_names=['run'],
order_by=['start_time desc'],
)
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_success(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = MagicMock(
experiment_id='0')
output = mlflow_repository.get_experiment('test')
assert output == 0
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_error(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = None
try:
mlflow_repository.get_experiment('test')
except ValueError as e:
assert str(e) == 'Experiment test not found'
else:
assert False
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run(mlflow, mlflow_repository):
mlflow.search_runs.return_value = DataFrame({
'params.retrain': ['True', 'False', 'True', 'False'],
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
'run_id': ['0', '1', '2', '3'],
})
output = mlflow_repository.get_experiment_last_run(0)
mlflow.search_runs.assert_called_once_with(
experiment_ids=[0],
filter_string="",
output_format="pandas",
)
assert output == '2'
@patch('laborious.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
client_mock = MagicMock()
mlflow.tracking.MlflowClient.return_value = client_mock
client_mock.get_registered_model.return_value = MagicMock(
latest_versions=[
MagicMock(version='1'),
MagicMock(version='2'),
MagicMock(version='3'),
]
)
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
mlflow.register_model.assert_called_once_with(
"runs:/0/prediction_model",
'test',
)
mlflow.tracking.MlflowClient.assert_called_once()
client_mock.get_registered_model.assert_called_once_with('test')
client_mock.transition_model_version_stage.assert_called_once_with(
name='test',
version='3',
stage='Production',
archive_existing_versions=True,
)
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
}
def test_update_production_model(mlflow_repository):
connector = mlflow_repository
with patch.object(connector, 'get_experiment',
return_value='0') as get_experiment:
with patch.object(connector, 'get_experiment_last_run',
return_value='2') as get_experiment_last_run:
with patch.object(connector, 'update_production_model_by_run_id',
return_value={'model_name': 'test', 'version': '3',
'mlflow_run_id': '0'}) as update_production_model_by_run_id:
output = connector.update_production_model('0', 'test')
get_experiment.assert_called_once_with('0')
get_experiment_last_run.assert_called_once_with('0')
update_production_model_by_run_id.assert_called_once_with(
'2', 'test')
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
'mlflow_experiment_id': '0',
}
def test_transform_success(mlflow_repository):
data = 'data'
model_name = 'model'
output = mlflow_repository.transform(model_name, data, 1)
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 1)
assert output == {
'success': True,
'content': mlflow_repository.model_serving.get_cached_transform.return_value.to_dict.return_value
}
def test_transform_error(mlflow_repository):
data = 'data'
model_name = 'model'
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception(
'error')
output = mlflow_repository.transform(model_name, data, 1)
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 1)
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
def test_predict_success(mlflow_repository):
data = 'data'
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict.return_value = np.array(
[2, 3]
)
output = mlflow_repository.predict(model_name, data, 1)
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 1)
assert output['success'] is True
assert output['content'] == {'prediction': {
0: 3}, 'response_time': ANY}
def test_predict_error(mlflow_repository):
data = 'data'
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict = MagicMock(
side_effect=Exception('error')
)
output = mlflow_repository.predict(model_name, data, 1)
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 1)
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}

View File

@@ -1,259 +0,0 @@
from unittest.mock import Mock, patch, MagicMock, ANY, call
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from pytest import fixture
from laborious.utils.repository.opc_repository import OpcRepository
from sientia_do.notifications.models import NotificationLevel
from datetime import datetime
@fixture
def mock_logger():
return Mock()
@fixture
def opc_repository(mock_logger):
return OpcRepository(
name="test_repo",
url="opc.tcp://localhost:4840",
logger=mock_logger,
notification_handler=Mock(),
reconnection_interval=60,
server_uri="urn:test:server",
cert_path="/path/to/cert.pem",
private_key_path="/path/to/key.pem",
server_cert_path="/path/to/server_cert.pem"
)
@fixture
def mock_client():
with patch('laborious.utils.repository.opc_repository.Client') as mock:
client_instance = MagicMock()
mock.return_value = client_instance
yield client_instance
def test_init(opc_repository):
assert opc_repository.name == "test_repo"
assert opc_repository.url == "opc.tcp://localhost:4840"
assert opc_repository.server_uri == "urn:test:server"
assert opc_repository.cert_path == "/path/to/cert.pem"
assert opc_repository.private_key_path == "/path/to/key.pem"
assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
assert opc_repository.reconnection_interval == 60
assert opc_repository.client is None
assert opc_repository.last_reconnection_time is None
assert opc_repository.error_count == 0
def test_set_security(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.set_security()
mock_client.application_uri = "urn:test:server"
mock_client.set_security.assert_called_once_with(
SecurityPolicyBasic256,
certificate="/path/to/cert.pem",
private_key="/path/to/key.pem",
server_certificate="/path/to/server_cert.pem"
)
assert mock_client.secure_channel_timeout == 10000000
assert mock_client.session_timeout == 10000000
def test_set_security_missing_certificates(opc_repository):
opc_repository.cert_path = None
opc_repository.private_key_path = None
try:
opc_repository.set_security()
except ValueError as e:
assert str(
e) == "Certificate and private key paths must be provided for secure connection."
def test_connect_with_security(opc_repository, mock_client):
opc_repository.try_connect = MagicMock()
opc_repository.connect()
opc_repository.try_connect.assert_called_once()
assert opc_repository.client == mock_client
def test_connect_without_security(opc_repository, mock_client):
opc_repository.cert_path = None
opc_repository.try_connect = MagicMock()
opc_repository.set_security = MagicMock()
opc_repository.connect()
opc_repository.try_connect.assert_called_once()
opc_repository.set_security.assert_not_called()
assert opc_repository.client == mock_client
def test_try_connect_sucess(opc_repository):
opc_repository.last_reconnection_time = None
opc_repository.client = MagicMock()
opc_repository.try_connect()
opc_repository.client.connect.assert_called_once()
assert opc_repository.last_reconnection_time is not None
def test_try_connect_fail(opc_repository):
opc_repository.last_reconnection_time = None
opc_repository.client = MagicMock()
opc_repository.client.connect.side_effect = Exception("Test error")
opc_repository.try_connect()
opc_repository.client.connect.assert_called_once()
assert opc_repository.last_reconnection_time is not None
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.name}",
message="Failed to connect to OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
def test_disconnect(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.disconnect()
mock_client.disconnect.assert_called_once()
assert opc_repository.client is None
def test_validate_connection_none_client(opc_repository):
opc_repository.client = None
opc_repository.connect = MagicMock()
response = opc_repository.validate_connection()
assert response
opc_repository.connect.assert_called_once()
def test_validate_connection_error_count_disconnect_error(opc_repository):
opc_repository.error_count = 6
opc_repository.client = MagicMock()
opc_repository.disconnect = MagicMock(side_effect=Exception("Test error"))
opc_repository.connect = MagicMock()
response = opc_repository.validate_connection()
assert response == opc_repository.connect.return_value
opc_repository.disconnect.assert_called_once()
opc_repository.connect.assert_called_once()
opc_repository.logger.error.assert_has_calls(
[
call("Failed to disconnect from OPC server: Test error"),
]
)
@patch('laborious.utils.repository.opc_repository.hasattr', return_value=True)
@patch('laborious.utils.repository.opc_repository.datetime',
MagicMock(now=MagicMock(return_value=datetime(2025, 1, 1, 0, 0, 0))))
def test_validate_connection_lost_not_time_to_reconect(_mock_datetime, opc_repository):
opc_repository.error_count = 0
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.try_connect = MagicMock()
response = opc_repository.validate_connection()
opc_repository.try_connect.assert_not_called()
assert response is False
@patch('laborious.utils.repository.opc_repository.hasattr', return_value=True)
@patch('laborious.utils.repository.opc_repository.datetime',
MagicMock(now=MagicMock(return_value=datetime(2025, 1, 1, 1, 0, 0))))
def test_validate_connection_lost_time_to_reconect(_mock_datetime, opc_repository):
opc_repository.error_count = 0
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.try_connect = MagicMock()
response = opc_repository.validate_connection()
opc_repository.try_connect.assert_called_once()
assert response == opc_repository.try_connect.return_value
def test_validate_connection_failed(opc_repository):
opc_repository.client = MagicMock()
opc_repository.error_count = 0
output = opc_repository.validate_connection()
assert output is True
def test_write_data_validate_connection_do_nothing(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = MagicMock()
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
def test_write_data_validate_connection_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=False)
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_not_called()
def test_write_data_get_node_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.client.get_node.side_effect = Exception("Test error")
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.name}",
message="Failed to get node from OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
assert opc_repository.error_count == 1
def test_write_data(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = mock_client
mock_node = MagicMock()
mock_client.get_node.return_value = mock_node
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_node.write_value.assert_called_once()
opc_repository.logger.info.assert_called_once_with(
"Writing 42.0 - <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

View File

@@ -1,133 +0,0 @@
from os import environ
from laborious.utils.connectors_config import (build_mlflow_config,
build_opc_config,
build_postgres_config)
def test_build_mlflow_config_with_env_vars():
# Arrange
environ['MLFLOW_HOST'] = 'http://test-host'
environ['MLFLOW_PORT'] = '8080'
environ['MLFLOW_USERNAME'] = 'test-user'
environ['MLFLOW_PASSWORD'] = 'test-pass'
# Act
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://test-host'
assert config['port'] == 8080
assert config['username'] == 'test-user'
assert config['password'] == 'test-pass'
def test_build_mlflow_config_with_defaults():
# Arrange
# Clear any existing env vars
environ.pop('MLFLOW_HOST', None)
environ.pop('MLFLOW_PORT', None)
environ.pop('MLFLOW_USERNAME', None)
environ.pop('MLFLOW_PASSWORD', None)
# Act
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://localhost'
assert config['port'] == 5080
assert config['username'] == 'aignosi'
assert config['password'] == 'aignosi'
def test_build_opc_config_with_env_vars():
# Arrange
environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}'
# Act
config = build_opc_config()
# Assert
assert config['opc']['name'] == 'test-opc'
assert config['opc']['url'] == 'opc.tcp://test:4840'
def test_build_opc_config_with_individual_env_vars():
# Arrange
environ.pop('OPC_CONFIG', None)
environ['OPC_NAME'] = 'test-name'
environ['OPC_URL'] = 'opc.tcp://test:4840'
environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840'
environ['OPC_RECONNECTION_INTERVAL'] = '300'
# Act
config = build_opc_config()
# Assert
assert config['opc']['name'] == 'test-name'
assert config['opc']['url'] == 'opc.tcp://test:4840'
assert config['opc']['server_uri'] == 'opc.tcp://test:4840'
assert config['opc']['reconnection_interval'] == 300
def test_build_opc_config_with_defaults():
# Arrange
environ.pop('OPC_CONFIG', None)
environ.pop('OPC_NAME', None)
environ.pop('OPC_URL', None)
environ.pop('OPC_SERVER_URI', None)
environ.pop('OPC_RECONNECTION_INTERVAL', None)
# Act
config = build_opc_config()
# Assert
assert config['opc']['name'] == 'opc'
assert config['opc']['url'] == 'opc.tcp://localhost:4840'
assert config['opc']['server_uri'] == 'opc.tcp://localhost:4840'
assert config['opc']['reconnection_interval'] == 120
def test_build_postgres_config_with_env_vars():
# Arrange
environ['POSTGRES_HOST'] = 'test-host'
environ['POSTGRES_PORT'] = '5433'
environ['POSTGRES_USER'] = 'test-user'
environ['POSTGRES_PASSWORD'] = 'test-pass'
environ['POSTGRES_DBNAME'] = 'test-db'
environ['POSTGRES_MIN_CONNECTIONS'] = '10'
environ['POSTGRES_MAX_CONNECTIONS'] = '30'
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'test-host'
assert config['port'] == 5433
assert config['user'] == 'test-user'
assert config['password'] == 'test-pass'
assert config['dbname'] == 'test-db'
assert config['min_connections'] == 10
assert config['max_connections'] == 30
def test_build_postgres_config_with_defaults():
# Arrange
environ.pop('POSTGRES_HOST', None)
environ.pop('POSTGRES_PORT', None)
environ.pop('POSTGRES_USER', None)
environ.pop('POSTGRES_PASSWORD', None)
environ.pop('POSTGRES_DBNAME', None)
environ.pop('POSTGRES_MIN_CONNECTIONS', None)
environ.pop('POSTGRES_MAX_CONNECTIONS', None)
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'localhost'
assert config['port'] == 5432
assert config['user'] == 'sientia'
assert config['password'] == 'sientia'
assert config['dbname'] == 'sientia'
assert config['min_connections'] == 5
assert config['max_connections'] == 20

View File

@@ -1,37 +0,0 @@
import os
from unittest.mock import patch
import logging
import pytest
from laborious.utils.logger import get_logger
@pytest.fixture
def mock_env_vars():
with patch.dict(os.environ, {}, clear=True):
yield
@pytest.mark.usefixtures("mock_env_vars")
@patch('laborious.utils.logger.logging.Formatter')
@patch('laborious.utils.logger.logging.StreamHandler')
def test_get_logger_defaults(mock_stream_handler, mock_formatter):
"""Test logger creation with default settings"""
# Mock the StreamHandler and Formatter
logger = get_logger('test_logger')
# Verify logger settings
assert logger.name == 'test_logger'
assert logger.level == logging.INFO
# Verify handler configuration
mock_stream_handler.return_value.setLevel.assert_called_once_with('INFO')
mock_stream_handler.return_value.setFormatter.assert_called_once()
# Verify formatter configuration
mock_formatter.assert_called_once_with(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
# Verify handler was added to logger
assert len(logger.handlers) == 1

View File

@@ -1,127 +0,0 @@
from unittest.mock import call, patch, AsyncMock, ANY
from pytest import mark, fixture
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
@fixture
def format_and_export_prediction():
return FormatAndExportPrediction()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
"path_flag": None,
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"opc_servers": ["test_server"],
"opc_output_config": {"test": "config"}
}
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_prediction,
{
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence']
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
"path_flag": "default",
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"opc_servers": ["test_server"],
"opc_output_config": {"test": "config"},
"comment": "test_comment"
}
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_default_prediction,
{
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment']
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1

View File

@@ -1,515 +0,0 @@
from unittest.mock import AsyncMock, patch, call, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
@fixture
def prediction_process():
return PredictionProcess()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=False)
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'},
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
# mlflow_response_gate (predict)
('continue', 0.95, "Error"),
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'path_flag': 'continue',
'data': 'predicted_data',
'prediction_confidence': 0.95,
'timestamp': '2024-01-01',
'model_id': 1,
'model_name': 'test_model_name',
'model_retention': '30',
'opc_output_config': input_data['opc_output_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'comment': 'Error'
}
)
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=True)
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('stop', 0.95, "Input data with bad quality"), # input_gate
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('repeat', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 4
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']},
retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(
side_effect=[False, False, True])
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 5
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(
side_effect=[False, False, False, True])
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95, "Error"), # mlflow_response_gate (predict)
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'STOP'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention
}, confidence, last_timestamp, ""
)
# Assert
assert result is True
workflow_mock.execute_local_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'repeat'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention
}, confidence, last_timestamp, ""
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.repeat_last_prediction,
{
'schema': schema,
'table_name': table_name,
'model_id': model
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'CONTINUE'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention,
'opc_output_config': {'test': 'config'}
}, confidence, last_timestamp, 'Prediction Process'
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'path_flag': path_flag,
'data': data,
'prediction_confidence': confidence,
'timestamp': last_timestamp,
'model_id': model,
'model_name': model_name,
'model_retention': model_retention,
'schema': schema,
'table_name': table_name,
'comment': 'Prediction Process',
'opc_output_config': {'test': 'config'}
}
)
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'unknown'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention,
'opc_output_config': {'test': 'config'}
}, confidence, last_timestamp, ""
)
# Assert
assert result is False
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called()

View File

@@ -1,81 +0,0 @@
from unittest.mock import AsyncMock, call, patch, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@fixture
def predictions_batch() -> PredictionsBatch:
return PredictionsBatch()
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_local_activity_method.return_value = {
'data': 'test_data'
}
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'query': 'SELECT * FROM test',
'schema': 'test_schema',
'table_name': 'test_table',
'opc_output_config': 'test_opc_output_config'
}
await predictions_batch.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
input_data['query'],
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
prediction_input = {
'data': {'data': 'test_data'},
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'model_retention': input_data.get('model_retention', 60),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {})
}
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', prediction_input)
])

View File

@@ -0,0 +1,116 @@
from unittest.mock import patch, MagicMock, ANY
from pytest import mark
from orchestrator.activities.activities import Activities
from orchestrator.activities.couchbase import Couchbase
from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.activities.slot_manager import SlotManager
@patch('orchestrator.activities.couchbase.Couchbase.__init__')
@patch('orchestrator.activities.temporal_manager.TemporalManager.__init__')
@patch('orchestrator.activities.slot_manager.SlotManager.__init__')
def test___init__(mock_slot_manager_init, mock_temporal_manager_init,
mock_couchbase_init):
couchbase_config = {
'connection_string': 'couchbase://localhost',
'username': 'admin',
'password': 'password'
}
redis_config = {
'host': 'localhost',
'port': 6379,
'username': 'admin',
'password': 'password'
}
temporal_client = MagicMock()
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
temporal_client=temporal_client,
couchbase_config=couchbase_config,
redis_config=redis_config,
logger=logger,
notification_handler=notification_handler
)
assert isinstance(activities, Activities)
assert isinstance(activities, Couchbase)
assert isinstance(activities, TemporalManager)
assert isinstance(activities, SlotManager)
mock_slot_manager_init.assert_called_once_with(
ANY,
host="localhost",
port=6379,
username="admin",
password="password",
logger=logger,
notification_handler=notification_handler
)
mock_couchbase_init.assert_called_once_with(
ANY,
connection_string=couchbase_config['connection_string'],
username=couchbase_config['username'],
password=couchbase_config['password'],
logger=logger,
notification_handler=notification_handler
)
mock_temporal_manager_init.assert_called_once_with(
ANY,
temporal_client=temporal_client,
logger=logger,
notification_handler=notification_handler
)
@mark.asyncio
@patch('orchestrator.activities.couchbase.Cluster')
async def test_prepare_activity(_mock_cluster):
couchbase_config = {
'connection_string': 'couchbase://localhost',
'username': 'admin',
'password': 'password'
}
redis_config = {
'host': 'localhost',
'port': 6379,
'username': 'admin',
'password': 'password'
}
temporal_client = MagicMock()
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
temporal_client=temporal_client,
couchbase_config=couchbase_config,
redis_config=redis_config,
logger=logger,
notification_handler=notification_handler
)
input_data = {
'workflow_name': 'test-workflow-name',
'schedule_name': 'test-schedule-name',
'model_name': 'test-model-name',
'model_id': 'test-model-id'
}
await activities.prepare_activity(input_data)
assert activities.notification_handler.base_notification.pipeline == input_data[
'workflow_name']
assert activities.notification_handler.base_notification.trigger == input_data[
'schedule_name']
assert activities.notification_handler.base_notification.model_name == input_data[
'model_name']
assert activities.notification_handler.base_notification.model_id == input_data[
'model_id']

View File

@@ -0,0 +1,56 @@
from unittest.mock import MagicMock, patch, ANY
from pytest import fixture, mark, raises
from sientia_do.notifications.models import NotificationLevel
from orchestrator.activities.couchbase import Couchbase
@fixture
@patch("orchestrator.activities.couchbase.Cluster")
def couchbase(_cluster_mock):
return Couchbase(
connection_string="couchbase://localhost",
username="admin",
password="password",
logger=MagicMock(),
notification_handler=MagicMock(),
)
@mark.asyncio
async def test_load_query_from_couchbase_success(couchbase):
couchbase.cluster.query.return_value.rows.return_value = [
{"id": "1", "name": "test"},
{"id": "2", "name": "test2"},
]
query = "SELECT * FROM bucket"
result = await couchbase.load_query_from_couchbase({
"query": query,
})
assert result == [
{"id": "1", "name": "test"},
{"id": "2", "name": "test2"},
]
couchbase.cluster.query.assert_called_once_with(query)
couchbase.notification_handler.build_and_send_notification.assert_not_called()
@mark.asyncio
async def test_load_query_from_couchbase_failure(couchbase):
couchbase.cluster.query.side_effect = Exception("Test error")
query = "SELECT * FROM bucket"
with raises(Exception):
await couchbase.load_query_from_couchbase({
"query": query,
})
couchbase.cluster.query.assert_called_once_with(query)
couchbase.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="COUCHBASE_LOAD_QUERY_ERROR",
message="Failed to execute couchbase query: Test error",
block="load_query_from_couchbase",
level=NotificationLevel.ERROR,
attachment_content=ANY,
)

View File

@@ -0,0 +1,57 @@
from unittest.mock import MagicMock, patch
from pytest import mark, fixture
from orchestrator.activities.slot_manager import SlotManager
@fixture
@patch("orchestrator.activities.slot_manager.Redis.__init__")
def slot_manager(_redis_mock):
slot_manager = SlotManager(
host="localhost",
port=6379,
username="admin",
password="password",
logger=MagicMock(),
notification_handler=MagicMock()
)
slot_manager.redis_client = MagicMock()
slot_manager.logger = MagicMock()
slot_manager.notification_handler = MagicMock()
return slot_manager
@mark.asyncio
async def test_load_opc_slots_no_slot_keys(slot_manager):
slot_manager.redis_client.keys.return_value = []
assert await slot_manager.load_opc_slots() == {}
@mark.asyncio
async def test_load_opc_slots(slot_manager):
slot_manager.redis_client.keys.return_value = [
b"slot:opc_tags:1", b"slot:opc_tags:2", b"slot:opc_tags:3"]
slot_manager.redis_client.mget.return_value = [
b"value1", "value2", None]
response = await slot_manager.load_opc_slots()
assert response == {
"slot:opc_tags:1": "value1",
"slot:opc_tags:2": "value2",
"slot:opc_tags:3": None
}
@mark.asyncio
async def test_load_active_ingestors(slot_manager):
slot_manager.redis_client.keys.return_value = [
b"heartbeat:ingestor:1", b"heartbeat:ingestor:2", b"heartbeat:ingestor:3"]
response = await slot_manager.load_active_ingestors()
assert response == ["heartbeat:ingestor:1",
"heartbeat:ingestor:2", "heartbeat:ingestor:3"]

View File

@@ -0,0 +1,80 @@
from unittest.mock import MagicMock, patch, AsyncMock
import base64
import json
from pytest import fixture, mark
from orchestrator.activities.temporal_manager import TemporalManager
@fixture
def temporal_manager():
return TemporalManager(
temporal_client=MagicMock(),
logger=MagicMock(),
notification_handler=MagicMock()
)
@mark.asyncio
@patch("orchestrator.activities.temporal_manager.MessageToDict",
return_value={"data": base64.b64encode(json.dumps({"test": "test"}).encode('utf-8'))})
async def test_load_schedule(_mock_message_to_dict, temporal_manager):
# Create async iterator mock
async def async_iter():
yield MagicMock(
id="test-schedule-id",
search_attributes={
"Orchestrated": ["true"]
}
)
yield MagicMock(
id="test-schedule-id-2",
search_attributes={
"Attr": ["false"]
}
)
yield MagicMock(
id="test-schedule-id-3",
search_attributes={
"Attr": ["false"]
}
)
temporal_manager.temporal_client.list_schedules = AsyncMock(
return_value=async_iter())
temporal_manager.temporal_client.get_schedule.return_value = MagicMock(
describe=AsyncMock(
return_value=MagicMock(
schedule=MagicMock(
action=MagicMock(
args=[
MagicMock(
data=base64.b64encode(json.dumps(
{"test": "test"}).encode('utf-8'))
)
]
)
)
)
)
)
temporal_manager.temporal_client.get_schedule.return_value.describe \
.return_value.schedule.spec = MagicMock(
intervals=[
MagicMock(
every=MagicMock(
seconds=60
)
)
]
)
response = await temporal_manager.load_schedule()
temporal_manager.temporal_client.list_schedules.assert_called_once()
assert response == {
"test-schedule-id": {
"frequency": 60,
"data": {"test": "test"},
"handle": temporal_manager.temporal_client.get_schedule.return_value
}
}

View File

@@ -0,0 +1,81 @@
from unittest.mock import AsyncMock, patch, ANY, call
from pytest import fixture, mark
from orchestrator.workflows.orchestrator import Orchestrator
from orchestrator.activities.activities import Activities
@fixture
def orchestrator():
return Orchestrator()
@mark.asyncio
@patch("orchestrator.workflows.orchestrator.workflow", new_callable=AsyncMock)
async def test_run(workflow_mock, orchestrator):
input_data = {
"pipelines_query": "SELECT * FROM bucket",
"opc_servers_query": "SELECT * FROM servers",
"schedule_name": "test-schedule-name",
}
await orchestrator.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
{
"workflow_name": "orchestrator",
"schedule_name": "test-schedule-name",
"model_name": "-",
"model_id": "-"
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_query_from_couchbase,
{
"query": input_data["pipelines_query"]
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_query_from_couchbase,
{
"query": input_data["opc_servers_query"]
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_schedule,
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_opc_slots,
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_active_ingestors,
retry_policy=ANY,
start_to_close_timeout=ANY
)
])