765 lines
30 KiB
Python
765 lines
30 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import threading
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
from alembic.config import Config
|
|
from alembic.migration import MigrationContext
|
|
from alembic.operations import Operations
|
|
from fastapi.testclient import TestClient
|
|
from pydantic import SecretStr, ValidationError
|
|
from sqlalchemy import create_engine, func, inspect, select, text
|
|
from sqlalchemy import event as sqlalchemy_event
|
|
from sqlalchemy.engine import make_url
|
|
from sqlalchemy.exc import IntegrityError, SQLAlchemyError
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
from modelforge_api.db import get_session
|
|
from modelforge_api.domain.alert_notifications import NotificationStatus
|
|
from modelforge_api.domain.observability import AlertState
|
|
from modelforge_api.main import app
|
|
from modelforge_api.persistence.models import (
|
|
AlertHistoryEvent,
|
|
AlertNotificationDelivery,
|
|
AlertRuleRevision,
|
|
Base,
|
|
OperationalAlert,
|
|
)
|
|
from modelforge_api.services.alert_notifications import (
|
|
_claim,
|
|
dispatch_pending,
|
|
enqueue_alert_transition,
|
|
)
|
|
from modelforge_api.services.observability import ObservabilityService
|
|
from modelforge_api.services.recovery_fingerprint import (
|
|
ALERT_NOTIFICATION_SCHEMA,
|
|
PORTFOLIO_SCHEMA,
|
|
FingerprintCompatibility,
|
|
_without_alert_notifications,
|
|
_without_lab_execution,
|
|
audit_fingerprint_differences_are_compatible,
|
|
fingerprint_compatibility,
|
|
fingerprint_session,
|
|
)
|
|
from modelforge_api.settings import Settings, get_settings
|
|
|
|
ADMIN_TOKEN = "alert-notification-test-admin" # noqa: S105 - synthetic test credential
|
|
HOOK_URL = "https://hooks.example.test/secret-route-value"
|
|
|
|
|
|
@pytest.fixture
|
|
def session() -> Session:
|
|
engine = create_engine(
|
|
"sqlite+pysqlite:///:memory:",
|
|
connect_args={"check_same_thread": False},
|
|
poolclass=StaticPool,
|
|
)
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as value:
|
|
yield value
|
|
engine.dispose()
|
|
|
|
|
|
def settings(**overrides: object) -> Settings:
|
|
values: dict[str, object] = {
|
|
"_env_file": None,
|
|
"env": "test",
|
|
"alert_notifications_enabled": True,
|
|
"alert_webhook_url": SecretStr(HOOK_URL),
|
|
"alert_webhook_allowed_hosts": "hooks.example.test",
|
|
"alert_notification_retry_base_seconds": 1,
|
|
"operator_api_key": SecretStr(ADMIN_TOKEN),
|
|
}
|
|
values.update(overrides)
|
|
return Settings(**values) # type: ignore[arg-type]
|
|
|
|
|
|
def make_alert_event(
|
|
session: Session, *, severity: str = "CRITICAL", state: str = "FIRING"
|
|
) -> tuple[OperationalAlert, AlertHistoryEvent]:
|
|
rule = AlertRuleRevision(
|
|
key=f"notification-test-{uuid.uuid4()}",
|
|
revision=1,
|
|
alert_type="BACKUP_STALE",
|
|
signal="backup_stale",
|
|
slo_policy_id=None,
|
|
condition={},
|
|
pending_seconds=0,
|
|
severity=severity,
|
|
labels={},
|
|
cooldown_seconds=0,
|
|
recovery_condition={},
|
|
active=True,
|
|
fingerprint=uuid.uuid4().hex,
|
|
created_by="notification-test",
|
|
)
|
|
session.add(rule)
|
|
session.flush()
|
|
alert = OperationalAlert(
|
|
rule_id=rule.id,
|
|
fingerprint=hashlib.sha256(str(uuid.uuid4()).encode()).hexdigest(),
|
|
alert_type="BACKUP_STALE",
|
|
severity=severity,
|
|
state=state,
|
|
source="backup_sets",
|
|
subject_type="backup",
|
|
subject_ref="latest",
|
|
summary="Latest verified backup is stale",
|
|
details={"secret": "must-not-be-in-webhook"},
|
|
first_seen_at=datetime.now(UTC),
|
|
last_seen_at=datetime.now(UTC),
|
|
)
|
|
session.add(alert)
|
|
session.flush()
|
|
event = AlertHistoryEvent(
|
|
alert_id=alert.id,
|
|
from_state=None,
|
|
to_state=state,
|
|
actor="observability-engine",
|
|
reason="Signal condition detected",
|
|
evidence={"secret": "must-not-be-in-webhook"},
|
|
occurred_at=datetime.now(UTC),
|
|
)
|
|
session.add(event)
|
|
session.flush()
|
|
return alert, event
|
|
|
|
|
|
def test_settings_require_https_exact_host_allowlist_and_standard_port() -> None:
|
|
with pytest.raises(ValidationError):
|
|
Settings(
|
|
_env_file=None,
|
|
alert_notifications_enabled=True,
|
|
alert_webhook_url=SecretStr("http://hooks.example.test/private"),
|
|
alert_webhook_allowed_hosts="hooks.example.test",
|
|
)
|
|
with pytest.raises(ValidationError):
|
|
Settings(
|
|
_env_file=None,
|
|
alert_notifications_enabled=True,
|
|
alert_webhook_url=SecretStr("https://evil.example.test/private"),
|
|
alert_webhook_allowed_hosts="hooks.example.test",
|
|
)
|
|
with pytest.raises(ValidationError):
|
|
Settings(
|
|
_env_file=None,
|
|
alert_notifications_enabled=True,
|
|
alert_webhook_url=SecretStr("https://hooks.example.test:8443/private"),
|
|
alert_webhook_allowed_hosts="hooks.example.test",
|
|
)
|
|
with pytest.raises(ValidationError):
|
|
Settings(
|
|
_env_file=None,
|
|
alert_notifications_enabled=True,
|
|
alert_webhook_url=SecretStr("https://hooks.example.test/private"),
|
|
alert_webhook_allowed_hosts="*.example.test",
|
|
)
|
|
|
|
|
|
def test_enqueue_is_transactional_filtered_and_unique_per_history_event(session: Session) -> None:
|
|
config = settings(alert_webhook_min_severity="WARNING")
|
|
alert, event = make_alert_event(session)
|
|
enqueue_alert_transition(session, config, alert, event)
|
|
session.flush()
|
|
row = session.scalar(select(AlertNotificationDelivery))
|
|
assert row is not None
|
|
assert row.idempotency_key == f"alert:{event.id}:default"
|
|
assert row.event_kind == "FIRING"
|
|
assert row.payload["summary"] == alert.summary
|
|
assert "details" not in row.payload
|
|
assert "must-not-be-in-webhook" not in json.dumps(row.payload)
|
|
|
|
enqueue_alert_transition(session, config, alert, event)
|
|
with pytest.raises(IntegrityError):
|
|
session.flush()
|
|
session.rollback()
|
|
assert session.scalar(select(func.count()).select_from(AlertNotificationDelivery)) == 0
|
|
|
|
|
|
def test_enqueue_only_notifies_firing_or_resolved_at_configured_severity(session: Session) -> None:
|
|
config = settings(alert_webhook_min_severity="WARNING")
|
|
alert, event = make_alert_event(session, severity="INFO", state="PENDING")
|
|
enqueue_alert_transition(session, config, alert, event)
|
|
assert session.scalar(select(func.count()).select_from(AlertNotificationDelivery)) == 0
|
|
|
|
event.to_state = "RESOLVED"
|
|
alert.severity = "CRITICAL"
|
|
enqueue_alert_transition(session, config, alert, event)
|
|
assert session.scalar(select(AlertNotificationDelivery)).event_kind == "RESOLVED" # type: ignore[union-attr]
|
|
|
|
|
|
def test_firing_transition_and_outbox_commit_in_one_transaction(session: Session) -> None:
|
|
config = settings()
|
|
alert, _initial_event = make_alert_event(session, state="PENDING")
|
|
now = datetime.now(UTC)
|
|
ObservabilityService(session, config)._transition(
|
|
alert, AlertState.FIRING, "observability-engine", "pending duration satisfied", now, {}
|
|
)
|
|
session.commit()
|
|
delivery = session.scalar(select(AlertNotificationDelivery))
|
|
assert delivery is not None
|
|
assert delivery.status == NotificationStatus.PENDING.value
|
|
assert delivery.event_kind == "FIRING"
|
|
assert delivery.alert_history_event_id is not None
|
|
|
|
session.delete(delivery)
|
|
session.commit()
|
|
second, _event = make_alert_event(session, state="PENDING")
|
|
ObservabilityService(session, config)._transition(
|
|
second, AlertState.FIRING, "observability-engine", "pending duration satisfied", now, {}
|
|
)
|
|
session.flush()
|
|
session.rollback()
|
|
assert session.scalar(select(func.count()).select_from(AlertNotificationDelivery)) == 0
|
|
|
|
|
|
def test_dispatch_retries_then_delivers_without_following_redirects_or_exposing_body(
|
|
session: Session,
|
|
) -> None:
|
|
config = settings()
|
|
alert, event = make_alert_event(session)
|
|
enqueue_alert_transition(session, config, alert, event)
|
|
session.commit()
|
|
requests: list[httpx.Request] = []
|
|
|
|
def retry(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
return httpx.Response(503, text="remote response contains sensitive diagnostic")
|
|
|
|
assert dispatch_pending(session, config, transport=httpx.MockTransport(retry)) == 1
|
|
row = session.scalar(select(AlertNotificationDelivery))
|
|
assert row is not None
|
|
assert row.status == NotificationStatus.RETRY_WAIT.value
|
|
assert row.attempt_count == 1
|
|
assert row.last_error_code == "REMOTE_RETRYABLE"
|
|
assert row.last_http_status == 503
|
|
assert HOOK_URL not in repr(row)
|
|
assert row.last_error_code not in "remote response contains sensitive diagnostic"
|
|
assert requests[0].url == httpx.URL(HOOK_URL)
|
|
assert requests[0].headers["idempotency-key"] == row.idempotency_key
|
|
assert requests[0].extensions.get("timeout") is not None
|
|
|
|
row.next_attempt_at = datetime.now(UTC) - timedelta(seconds=1)
|
|
session.commit()
|
|
assert dispatch_pending(
|
|
session,
|
|
config,
|
|
transport=httpx.MockTransport(lambda request: httpx.Response(200, text="ignored")),
|
|
) == 1
|
|
session.refresh(row)
|
|
assert row.status == NotificationStatus.DELIVERED.value
|
|
assert row.delivered_at is not None
|
|
assert row.last_error_code is None
|
|
|
|
|
|
def test_dispatch_rejects_redirect_and_permanent_client_errors(session: Session) -> None:
|
|
config = settings()
|
|
alert, event = make_alert_event(session)
|
|
enqueue_alert_transition(session, config, alert, event)
|
|
session.commit()
|
|
requests: list[httpx.Request] = []
|
|
|
|
def redirect(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
return httpx.Response(302, headers={"Location": "https://evil.example.test/"})
|
|
|
|
assert dispatch_pending(session, config, transport=httpx.MockTransport(redirect)) == 1
|
|
row = session.scalar(select(AlertNotificationDelivery))
|
|
assert row is not None
|
|
assert row.status == NotificationStatus.FAILED.value
|
|
assert row.last_error_code == "REDIRECT_REJECTED"
|
|
assert len(requests) == 1
|
|
|
|
|
|
def test_dispatch_permanent_failure_after_attempt_limit(session: Session) -> None:
|
|
config = settings(alert_notification_max_attempts=1)
|
|
alert, event = make_alert_event(session)
|
|
enqueue_alert_transition(session, config, alert, event)
|
|
session.commit()
|
|
assert dispatch_pending(
|
|
session,
|
|
config,
|
|
transport=httpx.MockTransport(lambda request: httpx.Response(400, text="private body")),
|
|
) == 1
|
|
row = session.scalar(select(AlertNotificationDelivery))
|
|
assert row is not None
|
|
assert row.status == NotificationStatus.FAILED.value
|
|
assert row.last_error_code == "REMOTE_REJECTED"
|
|
assert row.last_http_status == 400
|
|
|
|
|
|
def test_schema_0037_recovery_holds_pending_and_rejects_unrelated_delivery_changes(
|
|
session: Session,
|
|
) -> None:
|
|
config = settings()
|
|
alert, event = make_alert_event(session)
|
|
enqueue_alert_transition(session, config, alert, event)
|
|
pending = session.scalar(select(AlertNotificationDelivery))
|
|
assert pending is not None
|
|
pending.next_attempt_at = datetime.now(UTC) + timedelta(minutes=1)
|
|
pending.lease_expires_at = datetime.now(UTC) + timedelta(minutes=1)
|
|
pending.claim_token_hash = "a" * 64
|
|
delivered = AlertNotificationDelivery(
|
|
alert_history_event_id=None,
|
|
route_key="default",
|
|
idempotency_key="test:already-delivered:default",
|
|
event_kind="TEST",
|
|
severity=None,
|
|
payload={"event": "TEST"},
|
|
status="DELIVERED",
|
|
attempt_count=1,
|
|
last_attempt_at=datetime.now(UTC),
|
|
delivered_at=datetime.now(UTC),
|
|
)
|
|
session.add(delivered)
|
|
session.commit()
|
|
source = _without_lab_execution(fingerprint_session(session))
|
|
|
|
pending.status = NotificationStatus.HELD_FOR_REVIEW.value
|
|
pending.next_attempt_at = None
|
|
pending.lease_expires_at = None
|
|
pending.claim_token_hash = None
|
|
session.commit()
|
|
restored = fingerprint_session(session)
|
|
contract = fingerprint_compatibility(
|
|
source, restored, source_schema_revision=ALERT_NOTIFICATION_SCHEMA
|
|
)
|
|
assert contract.compatible
|
|
assert contract.mode.startswith(
|
|
"BENCHMARK_RUNS_UPGRADE|20260923_0037|LAB_EXECUTION_UPGRADE|20260923_0037|"
|
|
"ALERT_NOTIFICATION_STRICT|"
|
|
)
|
|
assert audit_fingerprint_differences_are_compatible(source, restored, contract)
|
|
|
|
delivered.status = NotificationStatus.HELD_FOR_REVIEW.value
|
|
session.commit()
|
|
corrupted = fingerprint_session(session)
|
|
rejected = fingerprint_compatibility(
|
|
source, corrupted, source_schema_revision=ALERT_NOTIFICATION_SCHEMA
|
|
)
|
|
assert not rejected.compatible
|
|
|
|
|
|
def test_restore_reconciliation_holds_pending_deliveries_and_requires_review(
|
|
session: Session, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from modelforge_api.services.recovery import RecoveryService
|
|
|
|
statements: list[str] = []
|
|
|
|
class RecoveryEngine:
|
|
def scalar(self, _target: object, sql: str) -> str:
|
|
statements.append(sql)
|
|
if "from alert_notification_deliveries" in sql and "select count(*)" in sql:
|
|
return "2"
|
|
if "from audit_events where action" in sql:
|
|
return "0"
|
|
if "string_agg" in sql:
|
|
return ""
|
|
return "0"
|
|
|
|
service = RecoveryService(session, Settings(_env_file=None, env="test"))
|
|
service.engine = RecoveryEngine() # type: ignore[assignment]
|
|
monkeypatch.setattr(service, "_verify_restored_audit_chain", lambda _target: {})
|
|
monkeypatch.setattr(service, "_append_restored_audit_marker", lambda *_args, **_kwargs: "a" * 64)
|
|
monkeypatch.setattr(service, "_reconcile_selected_comparisons", lambda _target: {"interrupted": 0, "claims_cleared": 0})
|
|
monkeypatch.setattr(service, "_reconcile_creative_allocations", lambda _target: 0)
|
|
monkeypatch.setattr(service.audit, "write", lambda *_args, **_kwargs: None)
|
|
record = SimpleNamespace(id=uuid.uuid4())
|
|
plan = SimpleNamespace(
|
|
database_destination="postgresql+psycopg://modelforge@restore.example.test/restored",
|
|
backup_set=SimpleNamespace(backup_id="notification-recovery"),
|
|
)
|
|
|
|
result = service._reconcile_restored(record, plan) # type: ignore[arg-type]
|
|
|
|
assert any(
|
|
"update alert_notification_deliveries set status='HELD_FOR_REVIEW'" in sql
|
|
and "next_attempt_at=null" in sql
|
|
and "lease_expires_at=null" in sql
|
|
and "claim_token_hash=null" in sql
|
|
and "where status in ('PENDING','DISPATCHING','RETRY_WAIT')" in sql
|
|
for sql in statements
|
|
)
|
|
assert result["notification_deliveries_held_for_review"] == 2
|
|
assert result["outcome"] == "RECOVERY_RECONCILIATION_REQUIRED"
|
|
|
|
|
|
def test_schema_0037_upgrade_allows_only_empty_outbox_and_exact_marker() -> None:
|
|
engine = create_engine("sqlite+pysqlite:///:memory:")
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as session:
|
|
restored = fingerprint_session(session)
|
|
source = _without_alert_notifications(_without_lab_execution(restored))
|
|
contract = fingerprint_compatibility(
|
|
source, restored, source_schema_revision=PORTFOLIO_SCHEMA
|
|
)
|
|
assert contract.compatible
|
|
assert contract.mode.startswith(
|
|
"BENCHMARK_RUNS_UPGRADE|20260923_0036|LAB_EXECUTION_UPGRADE|20260923_0036|"
|
|
)
|
|
assert "ALERT_NOTIFICATION_UPGRADE|PORTFOLIO_STRICT|" in contract.mode
|
|
assert audit_fingerprint_differences_are_compatible(source, restored, contract)
|
|
|
|
populated = json.loads(json.dumps(restored))
|
|
populated["tables"]["alert_notification_deliveries"]["row_count"] = 1
|
|
populated["tables"]["alert_notification_deliveries"]["digest"] = "f" * 64
|
|
rejected = fingerprint_compatibility(
|
|
source, populated, source_schema_revision=PORTFOLIO_SCHEMA
|
|
)
|
|
assert not rejected.compatible
|
|
|
|
unrelated = FingerprintCompatibility(
|
|
True,
|
|
"UNRELATED|" + contract.mode,
|
|
contract.allow_audit_checkpoint_introduction,
|
|
contract.reason,
|
|
)
|
|
assert not audit_fingerprint_differences_are_compatible(source, restored, unrelated)
|
|
engine.dispose()
|
|
|
|
|
|
def test_pre_0037_fingerprint_rejects_a_future_outbox_physical_table() -> None:
|
|
engine = create_engine("sqlite+pysqlite:///:memory:")
|
|
Base.metadata.create_all(engine)
|
|
with Session(engine) as session:
|
|
current = fingerprint_session(session)
|
|
source = _without_lab_execution(_without_alert_notifications(current))
|
|
source["physical_tables"] = sorted(
|
|
[*source["physical_tables"], "alert_notification_deliveries"]
|
|
)
|
|
rejected = fingerprint_compatibility(
|
|
source, source, source_schema_revision=PORTFOLIO_SCHEMA
|
|
)
|
|
assert not rejected.compatible
|
|
assert "table contract" in rejected.reason
|
|
engine.dispose()
|
|
|
|
|
|
def test_schema_0037_migration_creates_indexes_and_refuses_history_downgrade() -> None:
|
|
migration_path = (
|
|
Path(__file__).resolve().parents[1]
|
|
/ "alembic/versions/20260923_0037_alert_notification_outbox.py"
|
|
)
|
|
spec = importlib.util.spec_from_file_location("alert_notification_migration", migration_path)
|
|
assert spec is not None and spec.loader is not None
|
|
migration = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(migration)
|
|
|
|
engine = create_engine("sqlite+pysqlite:///:memory:")
|
|
with engine.begin() as connection, Operations.context(MigrationContext.configure(connection)):
|
|
migration.upgrade()
|
|
assert migration.down_revision == "20260923_0036"
|
|
assert "ix_alert_notification_dispatch_due" in {
|
|
item["name"] for item in inspect(connection).get_indexes("alert_notification_deliveries")
|
|
}
|
|
migration.downgrade()
|
|
migration.upgrade()
|
|
connection.execute(
|
|
text(
|
|
"insert into alert_notification_deliveries "
|
|
"(id, route_key, idempotency_key, event_kind, payload, status, attempt_count, created_at) "
|
|
"values (:id, 'default', 'test:migration:default', 'TEST', '{}', 'PENDING', 0, :created_at)"
|
|
),
|
|
{"id": str(uuid.uuid4()), "created_at": datetime.now(UTC)},
|
|
)
|
|
with pytest.raises(RuntimeError, match="history exists"):
|
|
migration.downgrade()
|
|
engine.dispose()
|
|
|
|
|
|
def test_admin_test_action_queues_safe_delivery_and_viewer_is_denied(session: Session) -> None:
|
|
config = settings()
|
|
app.dependency_overrides[get_session] = lambda: session
|
|
app.dependency_overrides[get_settings] = lambda: config
|
|
client = TestClient(app)
|
|
try:
|
|
assert client.post("/api/v1/admin/operations/notifications/test").status_code == 401
|
|
response = client.post(
|
|
"/api/v1/admin/operations/notifications/test",
|
|
headers={"X-ModelForge-Admin-Token": ADMIN_TOKEN},
|
|
)
|
|
assert response.status_code == 202
|
|
assert HOOK_URL not in response.text
|
|
assert "secret-route-value" not in response.text
|
|
delivery_id = response.json()["delivery"]["id"]
|
|
listed = client.get(
|
|
"/api/v1/admin/operations/notification-deliveries",
|
|
headers={"X-ModelForge-Admin-Token": ADMIN_TOKEN},
|
|
)
|
|
assert listed.status_code == 200
|
|
assert listed.json()[0]["id"] == delivery_id
|
|
assert "payload" not in listed.json()[0]
|
|
|
|
app.dependency_overrides[get_settings] = lambda: settings(
|
|
alert_notifications_enabled=False
|
|
)
|
|
unavailable = client.post(
|
|
"/api/v1/admin/operations/notifications/test",
|
|
headers={"X-ModelForge-Admin-Token": ADMIN_TOKEN},
|
|
)
|
|
assert unavailable.status_code == 503
|
|
assert unavailable.json()["error"]["code"] == "NOTIFICATION_ROUTE_UNAVAILABLE"
|
|
assert unavailable.json()["error"]["message"] == (
|
|
"The deployment-managed alert notification route is unavailable"
|
|
)
|
|
assert HOOK_URL not in unavailable.text
|
|
|
|
viewer_token = "mfop_" + "v" * 43
|
|
viewer_settings = settings(
|
|
operator_api_key=None,
|
|
operator_identities=SecretStr(
|
|
json.dumps(
|
|
[
|
|
{
|
|
"credential_id": "viewer-test",
|
|
"subject": "alert-viewer",
|
|
"role": "viewer",
|
|
"token_sha256": hashlib.sha256(viewer_token.encode()).hexdigest(),
|
|
"expires_at": (datetime.now(UTC) + timedelta(hours=1)).isoformat(),
|
|
}
|
|
]
|
|
)
|
|
),
|
|
)
|
|
app.dependency_overrides[get_settings] = lambda: viewer_settings
|
|
denied = client.post(
|
|
"/api/v1/admin/operations/notifications/test",
|
|
headers={"X-ModelForge-Admin-Token": viewer_token},
|
|
)
|
|
assert denied.status_code == 403
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
client.close()
|
|
|
|
|
|
def test_postgres_concurrent_claim_skips_the_row_held_by_another_worker() -> None:
|
|
admin_url = os.getenv("MODELFORGE_TEST_POSTGRES_URL")
|
|
if not admin_url:
|
|
pytest.skip("requires managed MODELFORGE_TEST_POSTGRES_URL")
|
|
owner_url = os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL")
|
|
runtime_url = os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL")
|
|
if not owner_url or not runtime_url:
|
|
pytest.fail("managed notification claim test requires owner and runtime URLs")
|
|
if admin_url.startswith("postgresql://"):
|
|
admin_url = admin_url.replace("postgresql://", "postgresql+psycopg://", 1)
|
|
|
|
database = f"test_alert_claim_{uuid.uuid4().hex}"
|
|
admin_engine = create_engine(admin_url, isolation_level="AUTOCOMMIT")
|
|
with admin_engine.connect() as connection:
|
|
connection.execute(text(f'CREATE DATABASE "{database}" OWNER modelforge'))
|
|
|
|
owner_database_url = make_url(owner_url).set(database=database)
|
|
runtime_engine = create_engine(make_url(runtime_url).set(database=database), pool_pre_ping=True)
|
|
migration_config = Config()
|
|
migration_config.set_main_option(
|
|
"script_location", str(Path(__file__).resolve().parents[1] / "alembic")
|
|
)
|
|
migration_config.set_main_option(
|
|
"sqlalchemy.url",
|
|
owner_database_url.render_as_string(hide_password=False).replace("%", "%%"),
|
|
)
|
|
first_selected = threading.Event()
|
|
release_first = threading.Event()
|
|
results: dict[str, int] = {}
|
|
failures: list[SQLAlchemyError | TimeoutError] = []
|
|
|
|
@sqlalchemy_event.listens_for(runtime_engine, "after_cursor_execute")
|
|
def pause_first_claim(
|
|
_connection: object,
|
|
_cursor: object,
|
|
statement: str,
|
|
_parameters: object,
|
|
_context: object,
|
|
_executemany: bool,
|
|
) -> None:
|
|
normalized = statement.lower()
|
|
if (
|
|
threading.current_thread().name == "notification-claim-a"
|
|
and "from alert_notification_deliveries" in normalized
|
|
and "for update" in normalized
|
|
):
|
|
first_selected.set()
|
|
if not release_first.wait(timeout=10):
|
|
raise TimeoutError("first notification claim did not receive release signal")
|
|
|
|
try:
|
|
import importlib
|
|
|
|
importlib.import_module("alembic.command").upgrade(migration_config, "head")
|
|
with Session(runtime_engine) as session:
|
|
session.add(
|
|
AlertNotificationDelivery(
|
|
idempotency_key=f"test:{uuid.uuid4()}:default",
|
|
route_key="default",
|
|
event_kind="TEST",
|
|
payload={"schema_version": 1, "event": "TEST"},
|
|
status="PENDING",
|
|
attempt_count=0,
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
config = settings(alert_notification_batch_size=1)
|
|
now = datetime.now(UTC)
|
|
|
|
def claim(name: str) -> None:
|
|
try:
|
|
with Session(runtime_engine) as session:
|
|
results[name] = len(_claim(session, config, now))
|
|
except (SQLAlchemyError, TimeoutError) as error: # retained until both workers unblock
|
|
failures.append(error)
|
|
|
|
first = threading.Thread(target=claim, args=("first",), name="notification-claim-a")
|
|
second = threading.Thread(target=claim, args=("second",), name="notification-claim-b")
|
|
first.start()
|
|
assert first_selected.wait(timeout=10), "first worker did not lock the due row"
|
|
second.start()
|
|
second.join(timeout=10)
|
|
release_first.set()
|
|
first.join(timeout=10)
|
|
assert not first.is_alive() and not second.is_alive()
|
|
assert failures == []
|
|
assert results == {"first": 1, "second": 0}
|
|
finally:
|
|
release_first.set()
|
|
runtime_engine.dispose()
|
|
with admin_engine.connect() as connection:
|
|
connection.execute(text(f'DROP DATABASE IF EXISTS "{database}" WITH (FORCE)'))
|
|
admin_engine.dispose()
|
|
|
|
|
|
def test_postgres_expired_lease_fences_a_stale_dispatch_result() -> None:
|
|
admin_url = os.getenv("MODELFORGE_TEST_POSTGRES_URL")
|
|
if not admin_url:
|
|
pytest.skip("requires managed MODELFORGE_TEST_POSTGRES_URL")
|
|
owner_url = os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL")
|
|
runtime_url = os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL")
|
|
if not owner_url or not runtime_url:
|
|
pytest.fail("managed notification lease test requires owner and runtime URLs")
|
|
if admin_url.startswith("postgresql://"):
|
|
admin_url = admin_url.replace("postgresql://", "postgresql+psycopg://", 1)
|
|
|
|
database = f"test_alert_lease_{uuid.uuid4().hex}"
|
|
admin_engine = create_engine(admin_url, isolation_level="AUTOCOMMIT")
|
|
with admin_engine.connect() as connection:
|
|
connection.execute(text(f'CREATE DATABASE "{database}" OWNER modelforge'))
|
|
|
|
owner_database_url = make_url(owner_url).set(database=database)
|
|
runtime_engine = create_engine(make_url(runtime_url).set(database=database), pool_pre_ping=True)
|
|
migration_config = Config()
|
|
migration_config.set_main_option(
|
|
"script_location", str(Path(__file__).resolve().parents[1] / "alembic")
|
|
)
|
|
migration_config.set_main_option(
|
|
"sqlalchemy.url",
|
|
owner_database_url.render_as_string(hide_password=False).replace("%", "%%"),
|
|
)
|
|
first_delivery_started = threading.Event()
|
|
release_first_delivery = threading.Event()
|
|
results: dict[str, int] = {}
|
|
failures: list[Exception] = []
|
|
workers: list[threading.Thread] = []
|
|
delivery_id = uuid.uuid4()
|
|
|
|
class DelayedFailureTransport(httpx.BaseTransport):
|
|
def handle_request(self, request: httpx.Request) -> httpx.Response:
|
|
first_delivery_started.set()
|
|
if not release_first_delivery.wait(timeout=10):
|
|
raise AssertionError("first delivery did not receive release signal")
|
|
return httpx.Response(503, request=request)
|
|
|
|
class SuccessfulTransport(httpx.BaseTransport):
|
|
def handle_request(self, request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(200, request=request)
|
|
|
|
def dispatch(
|
|
name: str,
|
|
claim_time: datetime,
|
|
transport: httpx.BaseTransport,
|
|
) -> None:
|
|
try:
|
|
with Session(runtime_engine) as session:
|
|
results[name] = dispatch_pending(
|
|
session, config, now=claim_time, transport=transport
|
|
)
|
|
except (AssertionError, RuntimeError, SQLAlchemyError, TimeoutError, ValueError) as error:
|
|
failures.append(error)
|
|
|
|
try:
|
|
import importlib
|
|
|
|
importlib.import_module("alembic.command").upgrade(migration_config, "head")
|
|
with Session(runtime_engine) as session:
|
|
session.add(
|
|
AlertNotificationDelivery(
|
|
id=delivery_id,
|
|
idempotency_key=f"test:{uuid.uuid4()}:default",
|
|
route_key="default",
|
|
event_kind="TEST",
|
|
payload={"schema_version": 1, "event": "TEST"},
|
|
status="PENDING",
|
|
attempt_count=0,
|
|
)
|
|
)
|
|
session.commit()
|
|
|
|
config = settings(alert_notification_batch_size=1, alert_notification_lease_seconds=15)
|
|
claim_time = datetime.now(UTC)
|
|
first = threading.Thread(
|
|
target=dispatch,
|
|
args=("first", claim_time, DelayedFailureTransport()),
|
|
name="notification-dispatch-stale",
|
|
)
|
|
second = threading.Thread(
|
|
target=dispatch,
|
|
args=(
|
|
"second",
|
|
claim_time + timedelta(seconds=config.alert_notification_lease_seconds + 1),
|
|
SuccessfulTransport(),
|
|
),
|
|
name="notification-dispatch-reclaimed",
|
|
)
|
|
first.start()
|
|
workers.append(first)
|
|
assert first_delivery_started.wait(timeout=10), "first worker did not begin delivery"
|
|
second.start()
|
|
workers.append(second)
|
|
second.join(timeout=10)
|
|
assert not second.is_alive(), "replacement worker did not finish"
|
|
assert results.get("second") == 1
|
|
release_first_delivery.set()
|
|
first.join(timeout=10)
|
|
assert not first.is_alive(), "stale worker did not finish"
|
|
assert failures == []
|
|
assert results == {"second": 1, "first": 1}
|
|
|
|
with Session(runtime_engine) as session:
|
|
delivery = session.get(AlertNotificationDelivery, delivery_id)
|
|
assert delivery is not None
|
|
assert delivery.status == "DELIVERED"
|
|
assert delivery.attempt_count == 2
|
|
assert delivery.last_http_status == 200
|
|
assert delivery.last_error_code is None
|
|
assert delivery.delivered_at is not None
|
|
assert delivery.lease_expires_at is None
|
|
assert delivery.claim_token_hash is None
|
|
finally:
|
|
release_first_delivery.set()
|
|
for worker in workers:
|
|
worker.join(timeout=10)
|
|
runtime_engine.dispose()
|
|
with admin_engine.connect() as connection:
|
|
connection.execute(text(f'DROP DATABASE IF EXISTS "{database}" WITH (FORCE)'))
|
|
admin_engine.dispose()
|