34dbc886f3
Update project configuration and dependencies - Added .mypy_cache and .cursor to .gitignore. - Changed asyncio_default_fixture_loop_scope and asyncio_default_test_loop_scope to "session" in pyproject.toml. - Updated e2e testing dependencies in requirements-dev.txt, replacing fakeredis and mongomock with pytest-httpserver. - Updated requirements.txt to use sientia_do instead of a specific git commit. - Modified sonar-project.properties to remove a file from coverage exclusions. - Enhanced E2E test fixtures in e2e/conftest.py for better container management. - Cleaned up e2e test files related to CoreScouter and PIWebAPIScouter workflows.
405 lines
11 KiB
Python
405 lines
11 KiB
Python
"""
|
|
Shared helpers for Scouter end-to-end tests.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
|
|
from pymongo import MongoClient
|
|
from redis import Redis
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.engine import Engine
|
|
from temporalio.client import Client
|
|
|
|
SCENARIO_INPUTS_DIR = Path(__file__).resolve().parent / 'scenario_inputs'
|
|
SCOUTER_TASK_QUEUE = 'scouter-test-queue'
|
|
|
|
_NOW_MARKER = re.compile(r'^@now(?:([+-])(\d+)([smhd]))?$')
|
|
|
|
|
|
def _resolve_timestamp_marker(value: str) -> datetime:
|
|
"""
|
|
Resolve @now relative timestamp markers to timezone-aware datetimes.
|
|
|
|
Args:
|
|
- value: Marker string such as @now, @now-1h, or @now+30m
|
|
|
|
Return:
|
|
Resolved datetime in UTC
|
|
"""
|
|
match = _NOW_MARKER.match(value.strip())
|
|
if not match:
|
|
raise ValueError(f'Invalid timestamp marker: {value}')
|
|
|
|
now = datetime.now(UTC)
|
|
if match.group(1) is None:
|
|
return now
|
|
|
|
sign, amount, unit = match.group(1), int(match.group(2)), match.group(3)
|
|
delta_kwargs = {'seconds': 0, 'minutes': 0, 'hours': 0, 'days': 0}
|
|
if unit == 's':
|
|
delta_kwargs['seconds'] = amount
|
|
elif unit == 'm':
|
|
delta_kwargs['minutes'] = amount
|
|
elif unit == 'h':
|
|
delta_kwargs['hours'] = amount
|
|
elif unit == 'd':
|
|
delta_kwargs['days'] = amount
|
|
|
|
delta = timedelta(**delta_kwargs)
|
|
return now - delta if sign == '-' else now + delta
|
|
|
|
|
|
def _resolve_payload(node: Any) -> Any:
|
|
"""
|
|
Recursively resolve @now markers inside JSON-loaded structures.
|
|
|
|
Args:
|
|
- node: JSON node (dict, list, or scalar)
|
|
|
|
Return:
|
|
Structure with markers replaced by datetimes or formatted strings
|
|
"""
|
|
if isinstance(node, dict):
|
|
return {key: _resolve_payload(value) for key, value in node.items()}
|
|
if isinstance(node, list):
|
|
return [_resolve_payload(item) for item in node]
|
|
if isinstance(node, str) and node.startswith('@now'):
|
|
resolved = _resolve_timestamp_marker(node)
|
|
return resolved.strftime('%Y-%m-%d %H:%M:%S.%f%z')
|
|
return node
|
|
|
|
|
|
def load_scenario_input(name: str, **overrides: Any) -> dict[str, Any]:
|
|
"""
|
|
Load a scenario JSON file and apply optional overrides.
|
|
|
|
Args:
|
|
- name: Scenario slug with or without .json suffix
|
|
- overrides: Top-level keys merged into the loaded document
|
|
|
|
Return:
|
|
Parsed scenario document with @now markers resolved
|
|
"""
|
|
slug = name.removesuffix('.json')
|
|
path = SCENARIO_INPUTS_DIR / f'{slug}.json'
|
|
with path.open(encoding='utf-8') as handle:
|
|
payload = json.load(handle)
|
|
resolved = _resolve_payload(payload)
|
|
if overrides:
|
|
resolved.update(overrides)
|
|
return resolved
|
|
|
|
|
|
def make_workflow_id(prefix: str) -> str:
|
|
"""
|
|
Build a unique Temporal workflow id for E2E runs.
|
|
|
|
Args:
|
|
- prefix: Human-readable prefix for the workflow id
|
|
|
|
Return:
|
|
Unique workflow id string
|
|
"""
|
|
return f'{prefix}-{uuid.uuid4().hex[:12]}'
|
|
|
|
|
|
async def start_and_await_workflow(
|
|
client: Client,
|
|
workflow_run: Any,
|
|
input_data: dict[str, Any],
|
|
workflow_id: str,
|
|
*,
|
|
task_queue: str = SCOUTER_TASK_QUEUE,
|
|
timeout: float = 300.0,
|
|
) -> Any:
|
|
"""
|
|
Start a workflow on the E2E task queue and await its result.
|
|
|
|
Args:
|
|
- client: Temporal client from WorkflowEnvironment
|
|
- workflow_run: Workflow run method (e.g. Scouter.run)
|
|
- input_data: Workflow input payload
|
|
- workflow_id: Unique workflow id
|
|
- task_queue: Task queue name
|
|
- timeout: Maximum seconds to wait for completion
|
|
|
|
Return:
|
|
Workflow result (None for Scouter-family workflows)
|
|
"""
|
|
return await client.execute_workflow(
|
|
workflow_run,
|
|
input_data,
|
|
id=workflow_id,
|
|
task_queue=task_queue,
|
|
execution_timeout=timedelta(seconds=timeout),
|
|
)
|
|
|
|
|
|
def _coerce_mongo_document(document: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Convert string timestamps in seed documents to BSON datetimes for Mongo filters.
|
|
|
|
Args:
|
|
- document: Raw document dict from scenario JSON
|
|
|
|
Return:
|
|
Document with inserted_at as datetime when present
|
|
"""
|
|
doc = dict(document)
|
|
inserted_at = doc.get('inserted_at')
|
|
if isinstance(inserted_at, str):
|
|
doc['inserted_at'] = datetime.strptime(inserted_at, DATETIME_FORMAT_MS_WITH_TZ).replace(
|
|
tzinfo=UTC
|
|
)
|
|
return doc
|
|
|
|
|
|
def seed_raw_collection(
|
|
mongo_uri: str,
|
|
database: str,
|
|
schedule_name: str,
|
|
documents: list[dict[str, Any]],
|
|
) -> None:
|
|
"""
|
|
Insert raw Mongo documents into raw_<schedule_name>.
|
|
|
|
Args:
|
|
- mongo_uri: MongoDB connection string
|
|
- database: Database name
|
|
- schedule_name: Schedule slug used in collection name
|
|
- documents: Documents to insert
|
|
"""
|
|
collection_name = f'raw_{schedule_name}'
|
|
client = MongoClient(mongo_uri)
|
|
try:
|
|
collection = client[database][collection_name]
|
|
if documents:
|
|
collection.insert_many([_coerce_mongo_document(doc) for doc in documents])
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def seed_last_data_timestamp(
|
|
redis_client: Redis,
|
|
workflow_name: str,
|
|
schedule_name: str,
|
|
value: str,
|
|
) -> None:
|
|
"""
|
|
Pre-seed last_data_timestamp Redis key the same way RedisRepository.set stores strings.
|
|
|
|
Args:
|
|
- redis_client: Connected Redis client
|
|
- workflow_name: Workflow name segment in the key
|
|
- schedule_name: Schedule name segment in the key
|
|
- value: Timestamp string to store
|
|
"""
|
|
key = f'last_data_timestamp:{workflow_name}:{schedule_name}'
|
|
redis_client.set(key, json.dumps(value))
|
|
|
|
|
|
def count_laborious_rows(engine: Engine, model_id: str | int | None = None) -> int:
|
|
"""
|
|
Count rows in sientia_data.laborious_data, optionally filtered by model_id.
|
|
|
|
Args:
|
|
- engine: SQLAlchemy engine bound to the Postgres testcontainer
|
|
- model_id: Optional model id filter
|
|
|
|
Return:
|
|
Row count
|
|
"""
|
|
query = 'SELECT COUNT(*) FROM sientia_data.laborious_data'
|
|
params: dict[str, Any] = {}
|
|
if model_id is not None:
|
|
query += ' WHERE model_id = :model_id'
|
|
params['model_id'] = int(model_id)
|
|
|
|
with engine.connect() as conn:
|
|
return conn.execute(text(query), params).scalar() or 0
|
|
|
|
|
|
def fetch_laborious_rows(engine: Engine, model_id: str | int | None = None) -> list[dict[str, Any]]:
|
|
"""
|
|
Fetch laborious_data rows as plain dicts.
|
|
|
|
Args:
|
|
- engine: SQLAlchemy engine bound to the Postgres testcontainer
|
|
- model_id: Optional model id filter
|
|
|
|
Return:
|
|
List of row dicts with variable and value keys
|
|
"""
|
|
query = 'SELECT model_id, variable, value, timestamp FROM sientia_data.laborious_data'
|
|
params: dict[str, Any] = {}
|
|
if model_id is not None:
|
|
query += ' WHERE model_id = :model_id'
|
|
params['model_id'] = int(model_id)
|
|
query += ' ORDER BY variable'
|
|
|
|
with engine.connect() as conn:
|
|
rows = conn.execute(text(query), params).mappings().all()
|
|
return [dict(row) for row in rows]
|
|
|
|
|
|
def count_held_data_keys(redis_client: Redis) -> int:
|
|
"""
|
|
Count Redis keys matching held_data_*.
|
|
|
|
Args:
|
|
- redis_client: Connected Redis client
|
|
|
|
Return:
|
|
Number of matching keys
|
|
"""
|
|
return len(redis_client.keys('held_data_*'))
|
|
|
|
|
|
def count_notifications(
|
|
mongo_uri: str,
|
|
database: str,
|
|
*,
|
|
notification_id: str | None = None,
|
|
level: str | None = None,
|
|
) -> int:
|
|
"""
|
|
Count notification documents in the E2E notification collection.
|
|
|
|
Args:
|
|
- mongo_uri: MongoDB connection string
|
|
- database: Database name
|
|
- notification_id: Optional notification_id filter
|
|
- level: Optional level filter (WARNING, ERROR, ...)
|
|
|
|
Return:
|
|
Matching document count
|
|
"""
|
|
client = MongoClient(mongo_uri)
|
|
try:
|
|
collection = client[database]['notification_queue']
|
|
query: dict[str, Any] = {}
|
|
if notification_id:
|
|
query['notification_id'] = notification_id
|
|
if level:
|
|
query['level'] = level
|
|
return collection.count_documents(query)
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def default_model_tags(
|
|
*,
|
|
names: list[str],
|
|
aggr: str = 'avg',
|
|
data_range: list[float] | None = None,
|
|
frequency: int = 60000,
|
|
) -> dict[str, dict[str, Any]]:
|
|
"""
|
|
Build a minimal model_tags map for E2E scenarios.
|
|
|
|
Args:
|
|
- names: Tag names to include
|
|
- aggr: Aggregation function (aggr_func field)
|
|
- data_range: Optional [min, max] validation range
|
|
- frequency: Collection frequency in milliseconds
|
|
|
|
Return:
|
|
model_tags dict keyed by tag name
|
|
"""
|
|
if data_range is None:
|
|
data_range = [0, 100]
|
|
return {
|
|
name: {
|
|
'webid': f'webid_{name}',
|
|
'aggr_func': aggr,
|
|
'data_range': data_range,
|
|
'frequency': frequency,
|
|
}
|
|
for name in names
|
|
}
|
|
|
|
|
|
def default_scouter_input(
|
|
*,
|
|
model_id: str = '1',
|
|
model_name: str = 'Test Model',
|
|
schedule_name: str = 'test-schedule',
|
|
workflow_name: str = 'scouter',
|
|
**overrides: Any,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Return a base workflow input dict for Scouter / CoreScouter E2E runs.
|
|
|
|
Args:
|
|
- model_id: Model identifier
|
|
- model_name: Human-readable model name
|
|
- schedule_name: Schedule slug
|
|
- workflow_name: Parent workflow name
|
|
- overrides: Additional keys merged into the payload
|
|
|
|
Return:
|
|
Workflow input dictionary
|
|
"""
|
|
payload: dict[str, Any] = {
|
|
'topic': 'e2e-topic',
|
|
'model_id': model_id,
|
|
'model_name': model_name,
|
|
'schedule_name': schedule_name,
|
|
'workflow_name': workflow_name,
|
|
'trigger_laborious': False,
|
|
'filters': {},
|
|
'schema': 'sientia_data',
|
|
'table_name': 'laborious_data',
|
|
'retention_time': 3600,
|
|
'fill_missing_tags': False,
|
|
'model_tags': default_model_tags(names=['tag1']),
|
|
}
|
|
payload.update(overrides)
|
|
return payload
|
|
|
|
|
|
def apply_pi_web_api_server_config(server: Any, config: dict[str, Any] | None) -> None:
|
|
"""
|
|
Configure the in-process PI Web API stub from a scenario pi_web_api_server block.
|
|
|
|
Args:
|
|
- server: PIWebAPITestServer instance
|
|
- config: Optional mode/rows/timeout_sleep_seconds dict from scenario JSON
|
|
"""
|
|
if not config:
|
|
return
|
|
server.set_mode(
|
|
config.get('mode', 'success'),
|
|
rows=config.get('rows'),
|
|
timeout_sleep_seconds=config.get('timeout_sleep_seconds', 60),
|
|
)
|
|
|
|
|
|
def postgres_connection_parts(connection_url: str) -> dict[str, Any]:
|
|
"""
|
|
Parse a SQLAlchemy Postgres URL into Activities postgres_config fields.
|
|
|
|
Args:
|
|
- connection_url: SQLAlchemy connection URL from testcontainers
|
|
|
|
Return:
|
|
Dict with host, port, user, password, dbname keys
|
|
"""
|
|
parsed = urlparse(connection_url)
|
|
return {
|
|
'host': parsed.hostname or 'localhost',
|
|
'port': parsed.port or 5432,
|
|
'user': parsed.username or 'test',
|
|
'password': parsed.password or 'test',
|
|
'dbname': (parsed.path or '/test').lstrip('/'),
|
|
}
|