83d3a4482c
SIENTIAPDE-1646 Enhance OPC integration and E2E testing framework - Updated .gitignore to exclude '.cursor/*' for better file management. - Added new marker in pyproject.toml for E2E tests using OPC. - Introduced async OPC server fixtures and improved connection handling in conftest.py. - Enhanced error handling in OPC activities and added metrics for session management in metrics.py. - Updated E2E tests to cover new OPC scenarios, including session errors and reconnect handling. - Refactored helper functions to improve prediction assertion logic in helpers.py. - Documented new OPC scenarios in scenarios.md for clarity on expected outcomes.
369 lines
13 KiB
Python
369 lines
13 KiB
Python
"""
|
|
Shared helpers for E2E tests (Temporal workflows + PostgreSQL).
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from datetime import datetime
|
|
from decimal import Decimal
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.engine import Engine
|
|
|
|
SCENARIO_INPUTS_DIR = Path(__file__).parent / 'scenario_inputs'
|
|
|
|
|
|
def _replace_template_values(payload: Any, model_id: int) -> Any:
|
|
"""
|
|
Replace string placeholders in scenario payloads with the concrete model id.
|
|
|
|
Args:
|
|
payload: JSON-like structure loaded from scenario input file.
|
|
model_id: Model id used to render template placeholders.
|
|
|
|
Return:
|
|
Any: Payload with ``{{MODEL_ID}}`` replaced where applicable.
|
|
"""
|
|
if isinstance(payload, dict):
|
|
return {key: _replace_template_values(value, model_id) for key, value in payload.items()}
|
|
if isinstance(payload, list):
|
|
return [_replace_template_values(item, model_id) for item in payload]
|
|
if isinstance(payload, str):
|
|
if payload == '{{MODEL_ID}}':
|
|
return model_id
|
|
return payload.replace('{{MODEL_ID}}', str(model_id))
|
|
return payload
|
|
|
|
|
|
def load_scenario_input(file_name: str, model_id: int | None = None) -> dict[str, Any]:
|
|
"""
|
|
Load a scenario input JSON from ``e2e/scenario_inputs``.
|
|
|
|
Args:
|
|
file_name: JSON file name inside ``e2e/scenario_inputs``.
|
|
model_id: Optional model id used to render ``{{MODEL_ID}}`` placeholders.
|
|
|
|
Return:
|
|
dict[str, Any]: Input payload ready to be passed to workflow/activity calls.
|
|
"""
|
|
file_path = SCENARIO_INPUTS_DIR / file_name
|
|
with file_path.open('r', encoding='utf-8') as f:
|
|
payload = json.load(f)
|
|
|
|
if model_id is not None:
|
|
return _replace_template_values(payload, model_id)
|
|
return payload
|
|
|
|
|
|
async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0):
|
|
"""
|
|
Start a workflow and wait for its result.
|
|
|
|
Args:
|
|
client: Temporal client from WorkflowEnvironment.
|
|
workflow_run: Workflow run method (e.g. PredictionsBatch.run).
|
|
input_data: Workflow input payload.
|
|
workflow_id: Unique workflow id.
|
|
timeout: Max seconds to wait for completion.
|
|
|
|
Return:
|
|
Workflow result value.
|
|
"""
|
|
handle = await client.start_workflow(
|
|
workflow_run,
|
|
input_data,
|
|
id=workflow_id,
|
|
task_queue='test-queue',
|
|
)
|
|
return await asyncio.wait_for(handle.result(), timeout=timeout)
|
|
|
|
|
|
DEFAULT_BATCH_TIMESTAMP = '2024-01-01 12:00:00+00:00'
|
|
DEFAULT_PREDICTION_HISTORY_TIMESTAMP = '2024-01-01 12:00:00+00:00'
|
|
|
|
|
|
def insert_sample_data(
|
|
postgres_engine: Engine,
|
|
model_id: int,
|
|
values: list[Any],
|
|
*,
|
|
data_timestamp: str = DEFAULT_BATCH_TIMESTAMP,
|
|
) -> None:
|
|
"""
|
|
Replace laborious_data rows for a model_id with one row per value (sensor_1..n).
|
|
|
|
Args:
|
|
postgres_engine: SQLAlchemy engine.
|
|
model_id: Model id column value.
|
|
values: Per-sensor values; use string 'NULL' for SQL NULL.
|
|
data_timestamp: Timestamp and created_at for every inserted row; drives
|
|
``last_timestamp`` on the MinIO/query payload (max row time).
|
|
"""
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
|
values_sql = []
|
|
for i, value in enumerate(values):
|
|
values_sql.append(f"""
|
|
({model_id}, 'sensor_{i + 1}', {value}, '{data_timestamp}', '{data_timestamp}')
|
|
""")
|
|
insert_sql = f"""
|
|
INSERT INTO sientia_data.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES
|
|
{', '.join(values_sql)}
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
|
|
|
|
def insert_sample_prediction(
|
|
postgres_engine: Engine,
|
|
model_id: int,
|
|
*,
|
|
prediction_timestamp: str = DEFAULT_PREDICTION_HISTORY_TIMESTAMP,
|
|
) -> tuple[int, Decimal, Decimal, str]:
|
|
"""
|
|
Insert a single historical prediction row for REPEAT scenarios.
|
|
|
|
Args:
|
|
postgres_engine: SQLAlchemy engine.
|
|
model_id: Model id.
|
|
prediction_timestamp: Row ``timestamp`` (unique with model_id in tests).
|
|
|
|
Return:
|
|
tuple: (model_id, prediction, prediction_confidence, prediction_status) for assertions.
|
|
"""
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
|
insert_sql = f"""
|
|
INSERT INTO sientia_data.predictions (
|
|
model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time
|
|
)
|
|
VALUES (
|
|
{model_id}, '{prediction_timestamp}', 10, 0, 'Good', '', 0.1
|
|
)
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
return (model_id, Decimal(10), Decimal(0), 'Good')
|
|
|
|
|
|
def workflow_failure_message_chain(exc: BaseException) -> list[str]:
|
|
"""
|
|
Collect ``str()`` / ``message`` from an exception and its ``__cause__`` chain.
|
|
|
|
Args:
|
|
exc: Root exception (e.g. from ``pytest.raises``).
|
|
|
|
Return:
|
|
list[str]: Messages from root to innermost cause.
|
|
"""
|
|
messages: list[str] = []
|
|
current: BaseException | None = exc
|
|
while current is not None:
|
|
messages.append(getattr(current, 'message', None) or str(current) or repr(current))
|
|
current = current.__cause__
|
|
return messages
|
|
|
|
|
|
def assert_postgres_unique_violation_in_chain(exc: BaseException) -> None:
|
|
"""
|
|
Assert the exception chain mentions Postgres unique-constraint violation.
|
|
|
|
Args:
|
|
exc: Workflow or activity error from Temporal.
|
|
|
|
Raises:
|
|
AssertionError: If no link in the chain looks like UniqueViolation.
|
|
"""
|
|
chain = ' | '.join(workflow_failure_message_chain(exc))
|
|
assert 'UniqueViolation' in chain or 'unique_model_id_timestamp' in chain, (
|
|
f'Expected unique constraint violation in error chain, got: {chain}'
|
|
)
|
|
|
|
|
|
def assert_prediction_row_count(postgres_engine: Engine, model_id: int, expected: int) -> None:
|
|
"""
|
|
Assert how many prediction rows exist for a model_id.
|
|
|
|
Args:
|
|
postgres_engine: SQLAlchemy engine.
|
|
model_id: Model id filter.
|
|
expected: Expected row count.
|
|
"""
|
|
with postgres_engine.connect() as conn:
|
|
n = conn.execute(
|
|
text('SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = :m'),
|
|
{'m': model_id},
|
|
).scalar()
|
|
assert n == expected, f'Expected {expected} prediction rows, got {n}'
|
|
|
|
|
|
def assert_prediction(
|
|
postgres_engine: Engine,
|
|
model_id: int,
|
|
prediction: float = 0.5,
|
|
prediction_confidence: int | Decimal = 0,
|
|
prediction_status: str = 'Good',
|
|
comments: str = '',
|
|
comments_contains: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Assert exactly one prediction row exists for model_id with expected columns.
|
|
|
|
Args:
|
|
postgres_engine: SQLAlchemy engine.
|
|
model_id: Expected model_id.
|
|
prediction: Expected prediction value.
|
|
prediction_confidence: Expected confidence (int or Decimal for numeric column).
|
|
prediction_status: Expected status string.
|
|
comments: Expected exact comments string (ignored when ``comments_contains`` is set).
|
|
comments_contains: When set, assert this substring appears in comments.
|
|
"""
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
|
|
f'FROM sientia_data.predictions WHERE model_id = {model_id} '
|
|
f'ORDER BY created_at ASC'
|
|
)
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}'
|
|
row = prediction_rows[0]
|
|
assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}'
|
|
assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), (
|
|
f'Expected prediction={prediction}, got {row[1]}'
|
|
)
|
|
assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal(
|
|
str(prediction_confidence)
|
|
), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
|
|
assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
|
|
actual_comments = row[4] or ''
|
|
if comments_contains is not None:
|
|
assert comments_contains in actual_comments, (
|
|
f"Expected comments to contain '{comments_contains}', got '{actual_comments}'"
|
|
)
|
|
else:
|
|
assert actual_comments == comments, f"Expected comments='{comments}', got '{actual_comments}'"
|
|
|
|
|
|
def assert_continue(
|
|
postgres_engine: Engine,
|
|
model_id: int,
|
|
prediction_confidence: Decimal = Decimal(2),
|
|
comments: str = 'Input data with bad quality',
|
|
) -> None:
|
|
"""Assert one default-style prediction row after CONTINUE gate path."""
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
|
|
f'FROM sientia_data.predictions WHERE model_id = {model_id}'
|
|
)
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 1, 'Expected one prediction record despite warnings'
|
|
row = prediction_rows[0]
|
|
assert row[1] == 0, f'Expected prediction=0, got {row[1]}'
|
|
assert row[2] == prediction_confidence, (
|
|
f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
|
|
)
|
|
assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}"
|
|
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
|
|
|
|
|
|
def assert_stop(postgres_engine: Engine, model_id: int) -> None:
|
|
"""Assert no prediction rows for model_id."""
|
|
import pytest
|
|
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(f'SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = {model_id}')
|
|
)
|
|
count = result_query.scalar()
|
|
assert count == 0, f'Expected no predictions, but found {count} records'
|
|
|
|
|
|
def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple) -> None:
|
|
"""
|
|
Assert two prediction rows for model_id both match last_prediction.
|
|
|
|
Rows are compared in created_at order for stability.
|
|
|
|
Args:
|
|
postgres_engine: SQLAlchemy engine.
|
|
model_id: Model id.
|
|
last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
|
|
"""
|
|
import pytest
|
|
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, prediction_status '
|
|
f'FROM sientia_data.predictions WHERE model_id = {model_id} '
|
|
f'ORDER BY created_at ASC'
|
|
)
|
|
)
|
|
prediction_rows = result_query.fetchall()
|
|
assert len(prediction_rows) == 2, 'Expected two prediction records'
|
|
assert prediction_rows[0] == last_prediction, (
|
|
f'Expected first row {last_prediction}, got {prediction_rows[0]}'
|
|
)
|
|
assert prediction_rows[1] == last_prediction, (
|
|
f'Expected second row {last_prediction}, got {prediction_rows[1]}'
|
|
)
|
|
|
|
|
|
def make_workflow_id(prefix: str) -> str:
|
|
"""Build a unique workflow id using a prefix and current timestamp."""
|
|
return f'{prefix}-{datetime.now().timestamp()}'
|
|
|
|
|
|
def insert_target_data_for_drift(
|
|
postgres_engine: Engine,
|
|
model_id: int,
|
|
timestamps: list[str],
|
|
variables_values: dict[str, list[float]],
|
|
) -> None:
|
|
"""
|
|
Insert one row per (timestamp, variable) pair into ``laborious_data``.
|
|
|
|
Used by drift scenarios that need wide-format input where the pivot keeps a
|
|
full row for every timestamp.
|
|
|
|
Args:
|
|
- postgres_engine: SQLAlchemy engine bound to the test container.
|
|
- model_id: Model id stamped on every row.
|
|
- timestamps: ISO-8601 strings used both as ``timestamp`` and ``created_at``.
|
|
- variables_values: Mapping of variable name to a list of values; each list
|
|
must be the same length as ``timestamps``.
|
|
"""
|
|
for var_name, values in variables_values.items():
|
|
if len(values) != len(timestamps):
|
|
raise ValueError(
|
|
f"Variable '{var_name}' has {len(values)} values but {len(timestamps)} timestamps"
|
|
)
|
|
|
|
rows_sql = []
|
|
for index, ts in enumerate(timestamps):
|
|
for var_name, values in variables_values.items():
|
|
rows_sql.append(
|
|
f"({model_id}, '{var_name}', {values[index]}, '{ts}', '{ts}')"
|
|
)
|
|
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(
|
|
text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}')
|
|
)
|
|
if rows_sql:
|
|
conn.execute(
|
|
text(
|
|
'INSERT INTO sientia_data.laborious_data '
|
|
'(model_id, variable, value, "timestamp", created_at) VALUES '
|
|
+ ', '.join(rows_sql)
|
|
)
|
|
)
|
|
|
|
|