cleveragents-core/features/steps/registry_thread_safety_steps.py

"""Step definitions for TDD Issue #10409 — get_provider_registry() thread-safety race condition.

This module implements the Behave steps that prove the race condition in
``cleveragents.providers.registry.get_provider_registry``. Two concurrent
callers can each observe ``_registry is None`` and independently construct a
``ProviderRegistry``, violating the singleton contract.

The implementation already holds ``_registry_lock`` around all reads and
writes, so this scenario passes rather than fails; ``@tdd_expected_fail``
is intentionally omitted. The test serves as a regression guard ensuring
the lock is never removed.
"""

from __future__ import annotations

import contextlib
import threading
from typing import Any
from unittest.mock import MagicMock, patch

from behave import given, then, when  # type: ignore[import-untyped]

import cleveragents.providers.registry as _registry_module
from cleveragents.providers.registry import (
    ProviderRegistry,
    get_provider_registry,
    reset_provider_registry,
)


def _make_settings() -> Any:
    """Return a minimal fake Settings object for ProviderRegistry construction."""
    settings = MagicMock()
    settings.openai_api_key = None
    settings.anthropic_api_key = None
    settings.google_api_key = None
    settings.gemini_api_key = None
    settings.azure_api_key = None
    settings.openrouter_api_key = None
    settings.cohere_api_key = None
    settings.groq_api_key = None
    settings.together_api_key = None
    settings.default_provider = None
    settings.default_model = None
    settings.azure_openai_endpoint = None
    settings.azure_openai_api_version = None
    settings.azure_openai_deployment = None
    settings.openrouter_organization = None
    return settings


@given("the global provider registry has been reset")
def step_reset_global_registry(context: Any) -> None:
    """Reset the global _registry singleton to None to simulate a fresh start."""
    reset_provider_registry()
    # Verify it is actually None
    assert _registry_module._registry is None, (
        "Expected _registry to be None after reset"
    )


@when("two threads call get_provider_registry() simultaneously")
def step_two_threads_call_get_provider_registry(context: Any) -> None:
    """Spawn two threads that call get_provider_registry() simultaneously.

    A threading.Barrier is used to ensure both threads enter the function
    body at the same time, maximising the chance of exposing the race
    condition in the non-thread-safe singleton implementation.

    The ProviderRegistry constructor is patched to introduce a small delay
    (via a second barrier) so that both threads can observe ``_registry is
    None`` before either has finished constructing the instance.
    """
    results: list[ProviderRegistry] = []
    construction_count: list[int] = [0]
    errors: list[Exception] = []

    # Barrier to synchronise both threads at the entry point of
    # get_provider_registry() so they both observe _registry is None.
    entry_barrier = threading.Barrier(2, timeout=10)

    # Barrier inside the ProviderRegistry constructor to ensure both threads
    # have passed the ``if _registry is None`` check before either assigns
    # the new instance to ``_registry``.
    construction_barrier = threading.Barrier(2, timeout=10)

    original_init = ProviderRegistry.__init__

    def _slow_init(self: ProviderRegistry, settings: Any = None) -> None:
        """Wrap __init__ to synchronise both threads mid-construction."""
        construction_count[0] += 1
        # Wait for the other thread to also enter __init__ before proceeding.
        # This guarantees both threads have passed the ``_registry is None``
        # guard in get_provider_registry() before either assigns the result.
        with contextlib.suppress(threading.BrokenBarrierError):
            construction_barrier.wait(timeout=10)
        if settings is None:
            settings = _make_settings()
        original_init(self, settings)

    def _thread_body() -> None:
        try:
            # Synchronise both threads at the entry of get_provider_registry()
            with contextlib.suppress(threading.BrokenBarrierError):
                entry_barrier.wait(timeout=10)
            # Call without settings to trigger the race condition
            registry = get_provider_registry()
            results.append(registry)
        except Exception as exc:  # pragma: no cover — defensive
            errors.append(exc)

    with (
        patch.object(ProviderRegistry, "__init__", _slow_init),
        patch(
            "cleveragents.providers.registry.get_settings",
            return_value=_make_settings(),
        ),
    ):
        t1 = threading.Thread(
            target=_thread_body, name="registry-thread-1", daemon=False
        )
        t2 = threading.Thread(
            target=_thread_body, name="registry-thread-2", daemon=False
        )
        t1.start()
        t2.start()
        t1.join(timeout=15)
        t2.join(timeout=15)

    if errors:  # pragma: no cover — defensive
        raise errors[0]

    context.registry_results = results
    context.registry_construction_count = construction_count[0]


@then("both threads should receive the identical registry instance")
def step_both_threads_same_instance(context: Any) -> None:
    """Assert that both threads received the exact same ProviderRegistry object.

    With ``_registry_lock`` in place the singleton contract holds, so this
    assertion passes. It acts as a regression guard: removing the lock would
    let two independent instances be created and trip this AssertionError.
    """
    results: list[ProviderRegistry] = context.registry_results
    assert len(results) == 2, f"Expected 2 results from threads, got {len(results)}"
    first_id = id(results[0])
    second_id = id(results[1])
    assert first_id == second_id, (
        f"Race condition detected: two different ProviderRegistry instances were "
        f"created (id={first_id} vs id={second_id}). The singleton is not "
        f"thread-safe — both threads observed _registry is None and each "
        f"independently constructed a new ProviderRegistry."
    )


@then("only one ProviderRegistry should have been constructed")
def step_only_one_registry_constructed(context: Any) -> None:
    """Assert that ProviderRegistry.__init__ was called exactly once.

    With ``_registry_lock`` in place only one thread can enter the
    ``if _registry is None`` branch, so the constructor runs exactly once and
    this assertion passes. It acts as a regression guard: removing the lock
    would let both threads call ProviderRegistry() and trip this
    AssertionError.
    """
    count: int = context.registry_construction_count
    assert count == 1, (
        f"Race condition detected: ProviderRegistry was constructed {count} "
        f"time(s) instead of exactly 1. Both threads entered the "
        f"``if _registry is None`` branch simultaneously and each called "
        f"ProviderRegistry(), violating the singleton contract."
    )