Code import - branch feature/SIENTIAPDE-1646
This commit is contained in:
446
e2e/conftest.py
Normal file
446
e2e/conftest.py
Normal file
@@ -0,0 +1,446 @@
|
||||
"""Pytest configuration and fixtures for orchestrator E2E tests."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from google.protobuf.duration_pb2 import Duration
|
||||
from pymongo import MongoClient
|
||||
from redis import Redis
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.temporal.worker.prepare_worker import build_queue_name
|
||||
from sqlalchemy import create_engine
|
||||
from temporalio.api.enums.v1 import IndexedValueType
|
||||
from temporalio.api.operatorservice.v1 import AddSearchAttributesRequest
|
||||
from temporalio.api.workflowservice.v1 import (
|
||||
DescribeNamespaceRequest,
|
||||
RegisterNamespaceRequest,
|
||||
)
|
||||
from temporalio.client import Client
|
||||
from temporalio.common import SearchAttributeKey
|
||||
from temporalio.service import RPCError, RPCStatusCode
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
from testcontainers.core.container import DockerContainer
|
||||
from testcontainers.postgres import PostgresContainer
|
||||
|
||||
from e2e.helpers import MONGO_COLLECTIONS, ORCHESTRATOR_TASK_QUEUE
|
||||
from e2e.smtp_test_server import SmtpTestServer
|
||||
from e2e.stub_workflows import STUB_WORKFLOW_CLASSES
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.activities.formatters import schedule_types
|
||||
from orchestrator.workflows.alerts import Alerts
|
||||
from orchestrator.workflows.orchestrator import Orchestrator
|
||||
from orchestrator.workflows.reports import Reports
|
||||
from orchestrator.workflows.subworkflows.load_notification_package import (
|
||||
LoadNotificationPackage,
|
||||
)
|
||||
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
|
||||
|
||||
DB_SCHEMA_SQL_PATH = Path(__file__).parent / 'db_schema.sql'
|
||||
E2E_DATABASE = 'orchestrator_test'
|
||||
E2E_RUNTIMES = ('legacy', 'gpu')
|
||||
MANAGED_NAMESPACES = ('scouter', 'laborious')
|
||||
|
||||
E2E_SEARCH_ATTRIBUTES = [
|
||||
SearchAttributeKey.for_keyword('model_id'),
|
||||
SearchAttributeKey.for_keyword('model_name'),
|
||||
SearchAttributeKey.for_keyword('orchestrated'),
|
||||
]
|
||||
|
||||
E2E_NAMESPACE_SEARCH_ATTRIBUTES = {
|
||||
'model_id': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
|
||||
'model_name': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
|
||||
'orchestrated': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
|
||||
}
|
||||
|
||||
|
||||
async def register_namespace_if_missing(env: WorkflowEnvironment, namespace: str) -> None:
|
||||
"""
|
||||
Register a Temporal namespace on the local dev server and wait until it is ready.
|
||||
|
||||
Args:
|
||||
env: Session WorkflowEnvironment from start_local().
|
||||
namespace: Namespace name to register.
|
||||
"""
|
||||
service = env.client.service_client
|
||||
try:
|
||||
await service.workflow_service.register_namespace(
|
||||
RegisterNamespaceRequest(
|
||||
namespace=namespace,
|
||||
workflow_execution_retention_period=Duration(seconds=86400),
|
||||
)
|
||||
)
|
||||
except RPCError as err:
|
||||
if err.status != RPCStatusCode.ALREADY_EXISTS:
|
||||
raise
|
||||
|
||||
deadline = time.monotonic() + 5.0
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
await service.workflow_service.describe_namespace(
|
||||
DescribeNamespaceRequest(namespace=namespace)
|
||||
)
|
||||
return
|
||||
except RPCError:
|
||||
await asyncio.sleep(0.1)
|
||||
raise TimeoutError(f'Namespace {namespace} not ready within 5s')
|
||||
|
||||
|
||||
async def ensure_namespace_search_attributes(
|
||||
env: WorkflowEnvironment, namespace: str
|
||||
) -> None:
|
||||
"""
|
||||
Register the orchestrator search attributes on a namespace, if missing.
|
||||
|
||||
The local Temporal dev server only registers search attributes on the default
|
||||
namespace at start time. Schedules created in additional namespaces fail with
|
||||
"no mapping defined for search attribute ..." unless we explicitly add the
|
||||
same attribute mappings to those namespaces via the operator service.
|
||||
|
||||
Args:
|
||||
env: Session WorkflowEnvironment from start_local().
|
||||
namespace: Namespace where attributes must be available.
|
||||
"""
|
||||
service = env.client.service_client
|
||||
try:
|
||||
await service.operator_service.add_search_attributes(
|
||||
AddSearchAttributesRequest(
|
||||
namespace=namespace,
|
||||
search_attributes=dict(E2E_NAMESPACE_SEARCH_ATTRIBUTES),
|
||||
)
|
||||
)
|
||||
except RPCError as err:
|
||||
if err.status != RPCStatusCode.ALREADY_EXISTS:
|
||||
raise
|
||||
|
||||
|
||||
def temporal_host_from_env(env: WorkflowEnvironment) -> str:
|
||||
"""Return target host:port for the in-process Temporal dev server."""
|
||||
return env.client.service_client.config.target_host
|
||||
|
||||
|
||||
def mongo_uri_from_container(mongo_container) -> str:
|
||||
"""Build a Mongo connection string for the testcontainer."""
|
||||
port = mongo_container.get_exposed_port(27017)
|
||||
return f'mongodb://localhost:{port}'
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
def postgres_container():
|
||||
"""PostgreSQL testcontainer used by all E2E tests."""
|
||||
postgres = PostgresContainer('postgres:15')
|
||||
postgres.start()
|
||||
yield postgres
|
||||
postgres.stop()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
def mongo_container():
|
||||
"""MongoDB testcontainer used by real CoreNotificationHandler."""
|
||||
mongo = DockerContainer('mongo:7').with_exposed_ports(27017)
|
||||
mongo.start()
|
||||
yield mongo
|
||||
mongo.stop()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
def redis_container():
|
||||
"""Redis testcontainer for slot and notification timestamp paths."""
|
||||
redis = DockerContainer('redis:7').with_exposed_ports(6379)
|
||||
redis.start()
|
||||
yield redis
|
||||
redis.stop()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def postgres_engine(postgres_container):
|
||||
"""SQLAlchemy engine bound to the PostgreSQL testcontainer."""
|
||||
engine = create_engine(postgres_container.get_connection_url())
|
||||
yield engine
|
||||
engine.dispose()
|
||||
|
||||
|
||||
def _create_schema_and_tables(engine):
|
||||
sql_text = DB_SCHEMA_SQL_PATH.read_text(encoding='utf-8')
|
||||
with engine.begin() as conn:
|
||||
conn.exec_driver_sql(sql_text)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True)
|
||||
def setup_postgres_schema_and_tables(postgres_engine):
|
||||
"""Recreate Postgres schema from e2e/db_schema.sql before each test."""
|
||||
_create_schema_and_tables(postgres_engine)
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mongo_uri(mongo_container):
|
||||
return mongo_uri_from_container(mongo_container)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True)
|
||||
def reset_mongo_collections(mongo_uri):
|
||||
"""Drop orchestrator-managed Mongo collections between tests."""
|
||||
client = MongoClient(mongo_uri)
|
||||
try:
|
||||
db = client[E2E_DATABASE]
|
||||
for name in MONGO_COLLECTIONS:
|
||||
db[name].drop()
|
||||
finally:
|
||||
client.close()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_client(redis_container):
|
||||
"""Redis client bound to the testcontainer."""
|
||||
port = int(redis_container.get_exposed_port(6379))
|
||||
client = Redis(host='localhost', port=port, decode_responses=True)
|
||||
yield client
|
||||
client.close()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True)
|
||||
def reset_redis(redis_client):
|
||||
"""Flush Redis between tests."""
|
||||
redis_client.flushdb()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def smtp_server():
|
||||
"""Session-scoped in-process SMTP server."""
|
||||
server = SmtpTestServer()
|
||||
server.start()
|
||||
yield server
|
||||
server.stop()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def smtp_messages_cleanup(smtp_server):
|
||||
"""Clear captured SMTP messages between tests."""
|
||||
smtp_server.clear()
|
||||
yield
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
async def temporal_env():
|
||||
"""Real Temporal dev server (schedule APIs supported)."""
|
||||
env = await WorkflowEnvironment.start_local(search_attributes=E2E_SEARCH_ATTRIBUTES)
|
||||
for namespace in MANAGED_NAMESPACES:
|
||||
await register_namespace_if_missing(env, namespace)
|
||||
await ensure_namespace_search_attributes(env, namespace)
|
||||
yield env
|
||||
await env.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def temporal_host(temporal_env):
|
||||
return temporal_host_from_env(temporal_env)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True)
|
||||
async def cleanup_temporal_schedules(temporal_env):
|
||||
"""Delete orphan schedules in scouter/laborious before each test."""
|
||||
host = temporal_host_from_env(temporal_env)
|
||||
for namespace in MANAGED_NAMESPACES:
|
||||
client = await Client.connect(host, namespace=namespace)
|
||||
async for schedule in await client.list_schedules():
|
||||
handle = client.get_schedule_handle(schedule.id)
|
||||
await handle.delete()
|
||||
yield
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def mock_logger():
|
||||
"""Logger double with readable console output for E2E runs."""
|
||||
logger = MagicMock(spec=Logger)
|
||||
logger.info = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
|
||||
logger.debug = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
|
||||
logger.error = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
|
||||
logger.warning = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
|
||||
logger.custom_info = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
|
||||
logger.custom_debug = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
|
||||
logger.custom_error = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
|
||||
logger.custom_warning = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
|
||||
return logger
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def notification_handler(mock_logger, mongo_container):
|
||||
"""Real notification handler using MongoDB testcontainer."""
|
||||
handler = CoreNotificationHandler(
|
||||
connection_string=mongo_uri_from_container(mongo_container),
|
||||
database=E2E_DATABASE,
|
||||
logger=mock_logger,
|
||||
project_name='orchestrator-e2e',
|
||||
)
|
||||
try:
|
||||
yield handler
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def notification_inserts(notification_handler):
|
||||
"""Spy on real Mongo insert calls issued by notification handler."""
|
||||
collection = notification_handler.mongo_collection
|
||||
original_insert_one = collection.insert_one
|
||||
spy = MagicMock(wraps=original_insert_one)
|
||||
collection.insert_one = spy
|
||||
try:
|
||||
yield spy
|
||||
finally:
|
||||
collection.insert_one = original_insert_one
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def test_activities(
|
||||
postgres_container,
|
||||
mongo_container,
|
||||
redis_container,
|
||||
smtp_server,
|
||||
mock_logger,
|
||||
notification_handler,
|
||||
temporal_host,
|
||||
):
|
||||
"""Real Activities wired to testcontainers and in-process SMTP."""
|
||||
mongo_port = mongo_container.get_exposed_port(27017)
|
||||
redis_port = int(redis_container.get_exposed_port(6379))
|
||||
|
||||
activities = Activities(
|
||||
temporal_config={
|
||||
'temporal_host': temporal_host,
|
||||
'temporal_scouter_namespace': 'scouter',
|
||||
'temporal_laborious_namespace': 'laborious',
|
||||
},
|
||||
redis_config={
|
||||
'host': 'localhost',
|
||||
'port': redis_port,
|
||||
'username': '',
|
||||
'password': '',
|
||||
},
|
||||
mongodb_config={
|
||||
'connection_string': f'mongodb://localhost:{mongo_port}',
|
||||
'database_name': E2E_DATABASE,
|
||||
'ttl_index_seconds': 3600,
|
||||
},
|
||||
email_config={
|
||||
'sender_email': 'e2e@example.com',
|
||||
'sender_password': '',
|
||||
'smtp_server': smtp_server.host,
|
||||
'smtp_port': smtp_server.port,
|
||||
},
|
||||
postgres_config={
|
||||
'host': 'localhost',
|
||||
'port': int(postgres_container.get_exposed_port(5432)),
|
||||
'user': postgres_container.username,
|
||||
'password': postgres_container.password,
|
||||
'dbname': postgres_container.dbname,
|
||||
'min_connections': 1,
|
||||
'max_connections': 5,
|
||||
},
|
||||
logger=mock_logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
await activities.connect_to_temporal()
|
||||
try:
|
||||
yield activities
|
||||
finally:
|
||||
activities.shutdown()
|
||||
|
||||
|
||||
def _orchestrator_activity_list(activities: Activities) -> list:
|
||||
return [
|
||||
activities.load_active_ingestors,
|
||||
activities.load_opc_slots,
|
||||
activities.update_slots,
|
||||
activities.delete_slots,
|
||||
activities.aggregate_documents_in_mongodb,
|
||||
activities.find_documents_in_mongodb,
|
||||
activities.update_pipelines_timestamps,
|
||||
activities.create_pipelines_timestamps,
|
||||
activities.delete_pipelines_timestamps,
|
||||
activities.create_collection_with_ttl_index,
|
||||
activities.create_schedules,
|
||||
activities.update_schedules,
|
||||
activities.delete_schedules,
|
||||
activities.normalize_schedules,
|
||||
activities.process_schedules,
|
||||
activities.process_slots,
|
||||
activities.create_schedule_config,
|
||||
activities.create_slot_config,
|
||||
activities.report_schedule_orchestration,
|
||||
activities.report_slot_orchestration,
|
||||
activities.format_schedule_config,
|
||||
activities.get_last_data_timestamp,
|
||||
activities.load_latest_data,
|
||||
activities.put_last_data_timestamp,
|
||||
activities.filter_notification_alerts,
|
||||
activities.filter_notification_reports,
|
||||
activities.build_email_html,
|
||||
activities.send_email,
|
||||
activities.format_log_report,
|
||||
activities.export_data_to_postgres,
|
||||
activities.store_notification_cache,
|
||||
]
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def orchestrator_worker(temporal_env, test_activities):
|
||||
"""Worker for orchestrator workflows and all activities on the default namespace."""
|
||||
with ThreadPoolExecutor(max_workers=32) as activity_executor:
|
||||
async with Worker(
|
||||
temporal_env.client,
|
||||
task_queue=ORCHESTRATOR_TASK_QUEUE,
|
||||
workflows=[
|
||||
Orchestrator,
|
||||
Alerts,
|
||||
Reports,
|
||||
LoadNotificationPackage,
|
||||
ProcessNotifications,
|
||||
],
|
||||
activities=_orchestrator_activity_list(test_activities),
|
||||
activity_executor=activity_executor,
|
||||
) as worker:
|
||||
yield worker
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def stub_workers(temporal_env):
|
||||
"""No-op workers on scouter/laborious namespaces for every managed workflow type."""
|
||||
host = temporal_host_from_env(temporal_env)
|
||||
worker_contexts: list[Worker] = []
|
||||
clients: list[Client] = []
|
||||
stub_types = list(schedule_types.keys()) + [
|
||||
'xgboost_predictions_batch',
|
||||
'xgboost_minimal_retrain',
|
||||
]
|
||||
|
||||
try:
|
||||
for namespace in MANAGED_NAMESPACES:
|
||||
client = await Client.connect(host, namespace=namespace)
|
||||
clients.append(client)
|
||||
queues = {
|
||||
build_queue_name(workflow_type, runtime)
|
||||
for workflow_type in stub_types
|
||||
for runtime in E2E_RUNTIMES
|
||||
}
|
||||
for queue in queues:
|
||||
worker = Worker(
|
||||
client,
|
||||
task_queue=queue,
|
||||
workflows=STUB_WORKFLOW_CLASSES,
|
||||
)
|
||||
await worker.__aenter__()
|
||||
worker_contexts.append(worker)
|
||||
yield worker_contexts
|
||||
finally:
|
||||
for worker in reversed(worker_contexts):
|
||||
await worker.__aexit__(None, None, None)
|
||||
Reference in New Issue
Block a user