133 lines
4.4 KiB
Python
133 lines
4.4 KiB
Python
"""Mock classes for unit and integration tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
|
|
class MockLLMProvider:
|
|
"""Mock LLM provider that returns pre-configured responses.
|
|
|
|
Usage::
|
|
|
|
provider = MockLLMProvider(responses=[{"key": "value"}, {"other": "result"}])
|
|
result = provider.extract_json("some prompt") # returns {"key": "value"}
|
|
result = provider.extract_json("another prompt") # returns {"other": "result"}
|
|
# After exhausting responses, returns last item repeatedly.
|
|
"""
|
|
|
|
def __init__(self, responses: list[Any] | None = None) -> None:
|
|
self._responses: list[Any] = responses if responses is not None else [{}]
|
|
self._call_count = 0
|
|
|
|
def extract_json(self, *args, **kwargs) -> Any:
|
|
"""Return the next configured response, cycling at the last one."""
|
|
idx = min(self._call_count, len(self._responses) - 1)
|
|
result = self._responses[idx]
|
|
self._call_count += 1
|
|
return result
|
|
|
|
def complete(self, *args, **kwargs) -> str:
|
|
"""Return the next configured response as a JSON string."""
|
|
return json.dumps(self.extract_json(*args, **kwargs))
|
|
|
|
@property
|
|
def call_count(self) -> int:
|
|
"""Number of times extract_json / complete was called."""
|
|
return self._call_count
|
|
|
|
def reset(self) -> None:
|
|
"""Reset the call counter."""
|
|
self._call_count = 0
|
|
|
|
|
|
class MockHTTPResponse:
|
|
"""Mock httpx-compatible HTTP response.
|
|
|
|
Mimics the interface used by httpx.Response / requests.Response so that
|
|
code that calls `.json()`, `.text`, `.status_code`, and
|
|
`.raise_for_status()` works without a real HTTP server.
|
|
|
|
Usage::
|
|
|
|
response = MockHTTPResponse(200, json_data={"id": 1}, text='{"id": 1}')
|
|
response.json() # {"id": 1}
|
|
response.raise_for_status() # no-op for 2xx
|
|
response.status_code # 200
|
|
|
|
error = MockHTTPResponse(404, json_data={"detail": "not found"})
|
|
error.raise_for_status() # raises RuntimeError
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
status_code: int = 200,
|
|
json_data: Any = None,
|
|
text: str = "",
|
|
) -> None:
|
|
self.status_code = status_code
|
|
self._json_data = json_data
|
|
self.text = text or (json.dumps(json_data) if json_data is not None else "")
|
|
|
|
def json(self) -> Any:
|
|
"""Return the configured JSON data."""
|
|
if self._json_data is None:
|
|
raise ValueError("No JSON data configured for this MockHTTPResponse")
|
|
return self._json_data
|
|
|
|
def raise_for_status(self) -> None:
|
|
"""Raise RuntimeError for 4xx/5xx status codes (mirrors httpx behaviour)."""
|
|
if self.status_code >= 400:
|
|
raise RuntimeError(
|
|
f"HTTP error {self.status_code}: {self.text}"
|
|
)
|
|
|
|
|
|
def mock_duckdb_connection(tables: dict[str, list[dict]] | None = None) -> MagicMock:
|
|
"""Return a MagicMock that mimics a DuckDB connection.
|
|
|
|
Args:
|
|
tables: Mapping of SQL pattern → list-of-tuples results that
|
|
``fetchall()`` should return when the executed SQL contains the
|
|
key as a substring. ``fetchone()`` returns the first tuple (or
|
|
None). If no key matches, fetchall returns [] and fetchone None.
|
|
|
|
The returned mock exposes:
|
|
- ``.execute(sql, params=None)`` — returns self (chainable)
|
|
- ``.fetchall()`` — returns matching rows or []
|
|
- ``.fetchone()`` — returns first matching row or None
|
|
- ``.close()`` — no-op
|
|
|
|
Example::
|
|
|
|
conn = mock_duckdb_connection({"SELECT * FROM users": [("alice", "admin")]})
|
|
conn.execute("SELECT * FROM users").fetchall() # [("alice", "admin")]
|
|
"""
|
|
tables = tables or {}
|
|
|
|
class _MockConn:
|
|
def __init__(self) -> None:
|
|
self._last_sql: str = ""
|
|
self._last_rows: list = []
|
|
|
|
def execute(self, sql: str, params: Any = None) -> "_MockConn":
|
|
self._last_sql = sql
|
|
self._last_rows = []
|
|
for pattern, rows in tables.items():
|
|
if pattern in sql:
|
|
self._last_rows = list(rows)
|
|
break
|
|
return self
|
|
|
|
def fetchall(self) -> list:
|
|
return self._last_rows
|
|
|
|
def fetchone(self) -> Any:
|
|
return self._last_rows[0] if self._last_rows else None
|
|
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
return _MockConn()
|