cleveragents-core/features/steps/audit_session_race_steps.py

"""Step definitions for bug #991 - AuditService._ensure_session() TOCTOU race.

This module provides the Behave step implementations for the TDD bug-capture
test that proves bug #991 is fixed.  The bug was a Time-Of-Check-to-Time-Of-Use
(TOCTOU) race in AuditService._ensure_session(): the method checked
if self._session is None and then created an engine, ran create_all,
built a sessionmaker, and assigned self._session - all without holding
a threading.Lock.  Concurrent callers could each see _session is None
and create duplicate engines, leaking all but the last.

The fix adds a threading.Lock with double-checked locking so only the
first thread to acquire the lock performs session initialisation.

Originally tagged @tdd_expected_fail so the expected assertion failure
(proving the bug exists) would pass CI.  After the fix for #991 was
implemented, @tdd_expected_fail was removed and the test now runs
normally - verifying that create_engine is called exactly once even
under concurrent access.

See: https://git.cleverthis.com/cleveragents/cleveragents-core/issues/991
"""

from __future__ import annotations

import contextlib
import os
import tempfile
import threading
from unittest.mock import MagicMock, patch

from behave import given, then, when
from behave.runner import Context
from sqlalchemy import create_engine as _real_create_engine
from sqlalchemy.engine import Engine

from cleveragents.application.services.audit_service import AuditService
from cleveragents.config.settings import Settings


def _make_settings(**overrides: object) -> Settings:
    """Create a Settings instance with optional field overrides.

    Resets the Settings singleton so a fresh instance is constructed.
    This follows the established project pattern used in environment.py
    and dozens of other step files.
    """
    Settings._instance = None
    base = Settings()
    if overrides:
        return base.model_copy(update=overrides)
    return base


@given("an AuditService with no pre-injected session and a temp database")
def step_audit_service_no_session(context: Context) -> None:
    """Create an AuditService that will use _ensure_session() for lazy init."""
    fd, db_path = tempfile.mkstemp(suffix=".db", prefix="test_race_991_")
    os.close(fd)
    os.unlink(db_path)
    context.db_path = db_path
    db_url = f"sqlite:///{db_path}"
    settings = _make_settings(database_url=db_url, audit_async=False)
    context.service = AuditService(settings=settings, database_url=db_url)
    context.thread_sessions = []
    context.thread_errors = []
    context.engine_create_count = 0

    def _cleanup_db(path: str = db_path) -> None:
        svc: AuditService = context.service
        if svc._session is not None:
            bind = svc._session.bind
            svc._session.close()
            if isinstance(bind, Engine):
                bind.dispose()
        for suffix in ("", "-journal", "-wal", "-shm"):
            with contextlib.suppress(OSError):
                os.unlink(path + suffix)

    context._cleanup_handlers.append(_cleanup_db)


@when("{n:d} threads call _ensure_session concurrently through a barrier")
def step_concurrent_ensure_session(context: Context, n: int) -> None:
    """Launch n threads that all call _ensure_session() simultaneously."""
    barrier = threading.Barrier(n)
    sessions: list[object] = []
    errors: list[Exception] = []
    collect_lock = threading.Lock()
    engine_mock = MagicMock(side_effect=_real_create_engine)

    def worker() -> None:
        """Thread worker: wait on barrier, then call _ensure_session()."""
        try:
            barrier.wait(timeout=10)
            session = context.service._ensure_session()
            with collect_lock:
                sessions.append(session)
        except Exception as exc:
            with collect_lock:
                errors.append(exc)

    with patch(
        "cleveragents.application.services.audit_service.create_engine",
        engine_mock,
    ):
        threads = [threading.Thread(target=worker, daemon=True) for _ in range(n)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=30)
        for t in threads:
            assert not t.is_alive(), (
                f"Thread {t.name} is still alive after join timeout - "
                f"possible deadlock in _ensure_session()"
            )
        barrier_errors = [
            err for err in errors if isinstance(err, threading.BrokenBarrierError)
        ]
        assert not barrier_errors, (
            "Barrier synchronization failed while coordinating concurrent "
            "_ensure_session() calls. This invalidates race setup. "
            f"Barrier errors: {barrier_errors}; all errors: {errors}"
        )
        context.engine_create_count = len(engine_mock.call_args_list)

    context.thread_sessions = sessions
    context.thread_errors = errors
    context.thread_count = n


@then("create_engine should have been called exactly once")
def step_create_engine_called_once(context: Context) -> None:
    """Assert create_engine was called exactly once."""
    count = context.engine_create_count
    assert count == 1, (
        f"create_engine was called {count} time(s), expected exactly 1. "
        f"This indicates the TOCTOU race in _ensure_session() is not fixed - "
        f"multiple threads saw _session as None because the threading.Lock is "
        f"missing or not effective. "
        f"Thread errors (if any): {context.thread_errors}"
    )


@then("all {n:d} threads should have received a valid session")
def step_all_threads_got_session(context: Context, n: int) -> None:
    """Verify that all n threads completed - either with a session or error."""
    total = len(context.thread_sessions) + len(context.thread_errors)
    assert total == n, (
        f"Expected {n} threads to complete but only {total} did "
        f"({len(context.thread_sessions)} succeeded, "
        f"{len(context.thread_errors)} errored). "
        "Each worker should report exactly one outcome (session or error). "
        "A mismatch indicates a hung thread or missing outcome capture."
    )