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
0
orchestrator/activities/__init__.py
Normal file
0
orchestrator/activities/__init__.py
Normal file
47
orchestrator/activities/activities.py
Normal file
47
orchestrator/activities/activities.py
Normal 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)
|
||||
94
orchestrator/activities/couchbase.py
Normal file
94
orchestrator/activities/couchbase.py
Normal 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
|
||||
130
orchestrator/activities/formatters.py
Normal file
130
orchestrator/activities/formatters.py
Normal 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
|
||||
}
|
||||
0
orchestrator/activities/notification.py
Normal file
0
orchestrator/activities/notification.py
Normal file
73
orchestrator/activities/slot_manager.py
Normal file
73
orchestrator/activities/slot_manager.py
Normal 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]
|
||||
66
orchestrator/activities/temporal_manager.py
Normal file
66
orchestrator/activities/temporal_manager.py
Normal 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
|
||||
Reference in New Issue
Block a user