LangBot/tests/factories/provider.py

"""
Fake provider factory for tests.

Provides a deterministic fake provider that simulates LLM responses without real API calls.
"""

from __future__ import annotations

from unittest.mock import Mock
import typing

import langbot_plugin.api.entities.builtin.provider.message as provider_message


class FakeProvider:
    """Deterministic fake provider for unit and integration tests.

    Simulates various provider behaviors:
    - Normal text response
    - Streaming response
    - Timeout error
    - Auth error
    - Rate-limit error
    - Malformed response

    Does not call real LLM vendors.
    Does not require API keys.
    """

    PONG_RESPONSE = "LANGBOT_FAKE_PONG"

    def __init__(
        self,
        *,
        default_response: str = "fake response",
        streaming_chunks: list[str] = None,
        raise_error: Exception = None,
        captured_requests: list = None,
    ):
        self._default_response = default_response
        self._streaming_chunks = streaming_chunks or ["fake ", "response"]
        self._raise_error = raise_error
        self._captured_requests = captured_requests if captured_requests is not None else []

    def returns(self, text: str) -> "FakeProvider":
        """Configure provider to return a specific text response."""
        self._default_response = text
        self._streaming_chunks = [text]
        return self

    def returns_streaming(self, chunks: list[str]) -> "FakeProvider":
        """Configure provider to return streaming chunks."""
        self._streaming_chunks = chunks
        self._default_response = "".join(chunks)
        return self

    def raises(self, error: Exception) -> "FakeProvider":
        """Configure provider to raise an error."""
        self._raise_error = error
        return self

    def timeout(self) -> "FakeProvider":
        """Configure provider to simulate timeout."""
        return self.raises(TimeoutError("Provider timeout"))

    def auth_error(self) -> "FakeProvider":
        """Configure provider to simulate auth error."""
        return self.raises(Exception("Invalid API key"))

    def rate_limit(self) -> "FakeProvider":
        """Configure provider to simulate rate limit."""
        return self.raises(Exception("Rate limit exceeded"))

    def malformed(self) -> "FakeProvider":
        """Configure provider to simulate malformed response."""
        self._default_response = None
        return self

    def get_captured_requests(self) -> list:
        """Get all captured request arguments for assertions."""
        return self._captured_requests.copy()

    def clear_captured_requests(self):
        """Clear captured requests."""
        self._captured_requests.clear()

    def _create_message(self, content: str) -> provider_message.Message:
        """Create a provider message from text content."""
        return provider_message.Message(
            role="assistant",
            content=content,
        )

    def _create_chunk(
        self,
        content: str,
        is_final: bool = False,
        msg_sequence: int = 0,
    ) -> provider_message.MessageChunk:
        """Create a provider message chunk."""
        return provider_message.MessageChunk(
            role="assistant",
            content=content,
            is_final=is_final,
            msg_sequence=msg_sequence,
        )

    async def invoke_llm(
        self,
        query,
        model,
        messages: list,
        funcs: list,
        extra_args: dict,
        remove_think: bool = False,
    ) -> provider_message.Message:
        """Simulate non-streaming LLM invocation."""
        # Capture request for assertions
        self._captured_requests.append({
            "query_id": query.query_id if query else None,
            "model": model.model_entity.name if model and hasattr(model, 'model_entity') else None,
            "messages": messages,
            "funcs": funcs,
            "extra_args": extra_args,
        })

        # Simulate error if configured
        if self._raise_error:
            raise self._raise_error

        # Return response
        if self._default_response is None:
            # Malformed response
            return provider_message.Message(role="assistant", content=None)

        return self._create_message(self._default_response)

    async def invoke_llm_stream(
        self,
        query,
        model,
        messages: list,
        funcs: list,
        extra_args: dict,
        remove_think: bool = False,
    ) -> typing.AsyncGenerator[provider_message.MessageChunk, None]:
        """Simulate streaming LLM invocation."""
        # Capture request for assertions
        self._captured_requests.append({
            "query_id": query.query_id if query else None,
            "model": model.model_entity.name if model and hasattr(model, 'model_entity') else None,
            "messages": messages,
            "funcs": funcs,
            "extra_args": extra_args,
            "streaming": True,
        })

        # Simulate error if configured
        if self._raise_error:
            raise self._raise_error

        # Yield chunks
        for i, chunk in enumerate(self._streaming_chunks):
            is_final = (i == len(self._streaming_chunks) - 1)
            yield self._create_chunk(chunk, is_final=is_final, msg_sequence=i)


def fake_provider(
    default_response: str = "fake response",
) -> FakeProvider:
    """Create a FakeProvider with optional default response."""
    return FakeProvider(default_response=default_response)


def fake_provider_pong() -> FakeProvider:
    """Create a FakeProvider that returns the pong response."""
    return FakeProvider(default_response=FakeProvider.PONG_RESPONSE)


def fake_provider_timeout() -> FakeProvider:
    """Create a FakeProvider that simulates timeout."""
    return FakeProvider().timeout()


def fake_provider_auth_error() -> FakeProvider:
    """Create a FakeProvider that simulates auth error."""
    return FakeProvider().auth_error()


def fake_provider_rate_limit() -> FakeProvider:
    """Create a FakeProvider that simulates rate limit."""
    return FakeProvider().rate_limit()


def fake_provider_malformed() -> FakeProvider:
    """Create a FakeProvider that simulates malformed response."""
    return FakeProvider().malformed()


# ============== Mock Model Factory ==============


def fake_model(
    *,
    uuid: str = "test-model-uuid",
    name: str = "test-model",
    abilities: list[str] = None,
    provider: FakeProvider = None,
) -> Mock:
    """Create a mock model with a fake provider."""
    model = Mock()
    model.model_entity = Mock()
    model.model_entity.uuid = uuid
    model.model_entity.name = name
    model.model_entity.abilities = abilities or ["func_call", "vision"]
    model.model_entity.extra_args = {}

    # Attach fake provider
    if provider is None:
        provider = FakeProvider()

    model.provider = provider

    return model