e6018af23f
Update README, requirements, and E2E tests for improved configuration and functionality - Enhanced the README with updated model configuration examples, including the addition of an alias for production. - Removed the `requirements-light.txt` file and updated `requirements-local.txt` and `requirements.txt` to replace `asyncua` with `opcua`. - Refactored E2E test scenarios to utilize scenario input files for better maintainability and clarity. - Improved test coverage for MinIO offload functionality and added new helper functions for loading scenario inputs. - Updated `values.yaml` to reflect new global configurations and environment variables for the laborious worker.
352 lines
13 KiB
Python
352 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
|
|
|
|
import pandas as pd
|
|
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)
|
|
|
|
|
|
def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]) -> 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.
|
|
"""
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM predictions_schema.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}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
|
""")
|
|
insert_sql = f"""
|
|
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
|
VALUES
|
|
{', '.join(values_sql)}
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
|
|
|
|
def assert_prediction(
|
|
postgres_engine: Engine,
|
|
model_id: int,
|
|
prediction: float = 0.5,
|
|
prediction_confidence: int | Decimal = 0,
|
|
prediction_status: str = 'Good',
|
|
comments: str = '',
|
|
) -> 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 comments string.
|
|
"""
|
|
import pytest
|
|
|
|
with postgres_engine.connect() as conn:
|
|
result_query = conn.execute(
|
|
text(
|
|
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
|
|
f'FROM predictions_schema.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]}"
|
|
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
|
|
|
|
|
|
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 predictions_schema.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 predictions_schema.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 predictions_schema.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 predictions_schema.laborious_data WHERE model_id = {model_id}')
|
|
)
|
|
if rows_sql:
|
|
conn.execute(
|
|
text(
|
|
'INSERT INTO predictions_schema.laborious_data '
|
|
'(model_id, variable, value, "timestamp", created_at) VALUES '
|
|
+ ', '.join(rows_sql)
|
|
)
|
|
)
|
|
|
|
|
|
def build_drift_dataframe(
|
|
timestamps: list[str],
|
|
features: list[str],
|
|
methods: list[str],
|
|
statistic: float = 1.0,
|
|
drift_flags: dict[tuple[str, str], bool] | None = None,
|
|
include_multivariate: bool = True,
|
|
multivariate_value: float = 16.0,
|
|
multivariate_drift: bool = True,
|
|
extra_rows: list[dict[str, Any]] | None = None,
|
|
) -> Any:
|
|
"""
|
|
Build a deterministic dataframe matching the schema returned by
|
|
``sientia_model.analytics.model_analysis.ModelAnalysis.get_drift_metrics_dataframe``
|
|
so drift E2E tests can pin the exact rows persisted to PostgreSQL.
|
|
|
|
The output mirrors the analyzer's canonical schema:
|
|
``timestamp, feature, metric, statistic, p_value, alert, chunk_index,
|
|
chunk_start_date, chunk_end_date``. ``calculate_drift`` then renames
|
|
``alert -> drift``, ``chunk_index -> chunk``, ``chunk_end_date ->
|
|
timestamp_end`` and drops ``p_value`` / ``chunk_start_date``.
|
|
|
|
Args:
|
|
- timestamps: Truncated chunk start timestamps (``'2024-01-01 12:00'`` for ``min``).
|
|
- features: Univariate feature names (one row per feature/method/timestamp).
|
|
- methods: Univariate methods such as ``kolmogorov_smirnov``.
|
|
- statistic: Default univariate statistic value.
|
|
- drift_flags: Optional override of the ``alert`` flag per ``(feature, method)`` pair.
|
|
- include_multivariate: Whether to add a final multivariate row block.
|
|
- multivariate_value: Value placed on multivariate rows.
|
|
- multivariate_drift: Drift flag placed on multivariate rows.
|
|
- extra_rows: Additional pre-built rows to append (used for dedup/p_value tests).
|
|
|
|
Return:
|
|
pandas.DataFrame with columns: timestamp, feature, metric, statistic,
|
|
p_value, alert, chunk_index, chunk_start_date, chunk_end_date.
|
|
"""
|
|
rows: list[dict[str, Any]] = []
|
|
drift_flags = drift_flags or {}
|
|
# Synthetic chunk-end offset that mirrors the high-precision boundary
|
|
# (``...:59.999999999``) emitted by ``ModelAnalysis`` for minute chunks.
|
|
# Computing via ``Timedelta`` instead of string concatenation keeps the
|
|
# helper safe for both minute- and second-precision timestamps.
|
|
chunk_span = pd.Timedelta(seconds=59, nanoseconds=999999999)
|
|
|
|
for chunk_index, ts in enumerate(timestamps):
|
|
chunk_start = pd.Timestamp(ts)
|
|
chunk_end = chunk_start + chunk_span
|
|
for feature in features:
|
|
for method in methods:
|
|
rows.append(
|
|
{
|
|
'timestamp': chunk_start,
|
|
'feature': feature,
|
|
'metric': method,
|
|
'statistic': statistic,
|
|
'p_value': 0.5,
|
|
'alert': drift_flags.get((feature, method), False),
|
|
'chunk_index': chunk_index,
|
|
'chunk_start_date': chunk_start,
|
|
'chunk_end_date': chunk_end,
|
|
}
|
|
)
|
|
if include_multivariate:
|
|
rows.append(
|
|
{
|
|
'timestamp': chunk_start,
|
|
'feature': 'multivariate',
|
|
'metric': 'multivariate',
|
|
'statistic': multivariate_value,
|
|
'p_value': 0.0,
|
|
'alert': multivariate_drift,
|
|
'chunk_index': chunk_index,
|
|
'chunk_start_date': chunk_start,
|
|
'chunk_end_date': chunk_end,
|
|
}
|
|
)
|
|
|
|
if extra_rows:
|
|
rows.extend(extra_rows)
|
|
|
|
return pd.DataFrame(rows)
|