test(proxy/utils): pin PrismaClient get_data behavior (subset of #29488)

Scaffolding dependency for the #29983 and #30327 regression tests, which
modify test_prisma_client_get_data.py. Only conftest.py and
test_prisma_client_get_data.py are taken from #29488; the other twelve test
files in that PR pin behavior unrelated to this backport and are dropped.
test_prisma_client_get_data.py and conftest.py are self-contained (no imports
from the dropped siblings) and pass on this line.

(cherry picked from commit 457f65eff9)

tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py

@@ -0,0 +1,387 @@
+"""Shared fixtures for tests/test_litellm/proxy/utils/prisma_and_spend/.
+
+All fixtures used by PR2 test files live here. Do NOT add fixtures inside
+individual test files; if a fixture is missing, add it here and update the
+Notion plan.
+
+The PrismaClient is exercised against a fully-mocked Prisma stack: the
+``prisma.Prisma`` constructor and the writer/reader wrappers are patched
+before PrismaClient.__init__ runs so the init code paths execute without
+needing a generated Prisma client or a real database.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from dataclasses import dataclass, field
+from email.message import EmailMessage
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterator, List, Optional
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[5]))
+
+
+VOLATILE_KEYS = frozenset(
+    {
+        "created_at",
+        "updated_at",
+        "checked_at",
+        "started_at",
+        "request_id",
+        "id",
+        "token",
+        "expires",
+        "expires_at",
+        "litellm_call_id",
+        "created",
+        "spend",
+        "last_refreshed_at",
+        "startTime",
+        "endTime",
+        "salt",
+    }
+)
+
+
+def normalize(data: Any, volatile: frozenset = VOLATILE_KEYS) -> Any:
+    """Recursively replace values for volatile keys with '<VOLATILE>'."""
+    if isinstance(data, dict):
+        return {
+            k: ("<VOLATILE>" if k in volatile else normalize(v, volatile))
+            for k, v in data.items()
+        }
+    if isinstance(data, list):
+        return [normalize(v, volatile) for v in data]
+    return data
+
+
+_PRISMA_TABLES: List[str] = [
+    "litellm_verificationtoken",
+    "litellm_teamtable",
+    "litellm_usertable",
+    "litellm_endusertable",
+    "litellm_organizationtable",
+    "litellm_proxymodeltable",
+    "litellm_modeltable",
+    "litellm_budgettable",
+    "litellm_spendlogs",
+    "litellm_config",
+    "litellm_usernotifications",
+    "litellm_healthchecktable",
+    "litellm_dailyuserspend",
+    "litellm_dailyteamspend",
+    "litellm_dailytagspend",
+    "litellm_managed_object_table",
+    "litellm_credentialstable",
+    "litellm_mcpservertable",
+    "litellm_audit_log",
+    "litellm_invitationlink",
+    "litellm_session_token_table",
+    "litellm_passthrough_endpoint_table",
+    "litellm_cron_job",
+    "litellm_passthrough_logs",
+    "litellm_promptstable",
+    "litellm_guardrailstable",
+    "litellm_managed_files",
+    "litellm_mcpusercredentials",
+    "litellm_objectpermissiontable",
+    "litellm_organizationmembership",
+]
+
+
+def _make_table_mock() -> MagicMock:
+    table = MagicMock()
+    table.find_unique = AsyncMock(return_value=None)
+    table.find_many = AsyncMock(return_value=[])
+    table.find_first = AsyncMock(return_value=None)
+    table.create = AsyncMock()
+    table.create_many = AsyncMock()
+    table.update = AsyncMock()
+    table.update_many = AsyncMock()
+    table.upsert = AsyncMock()
+    table.delete = AsyncMock()
+    table.delete_many = AsyncMock()
+    table.count = AsyncMock(return_value=0)
+    table.group_by = AsyncMock(return_value=[])
+    table.aggregate = AsyncMock(return_value={})
+    return table
+
+
+@pytest.fixture
+def mock_prisma_client() -> MagicMock:
+    """Bare ``db`` mock with all common LiteLLM_* tables stubbed.
+
+    Override individual return values in a test::
+
+        mock_prisma_client.db.litellm_usertable.find_unique.return_value = user
+    """
+    client = MagicMock(name="MockPrismaClient")
+    client.db = MagicMock(name="MockPrismaDB")
+    client.connect = AsyncMock()
+    client.disconnect = AsyncMock()
+    client.health_check = AsyncMock(return_value=[{"?column?": 1}])
+    client.proxy_logging_obj = MagicMock()
+    client.proxy_logging_obj.failure_handler = AsyncMock()
+    client.spend_log_transactions = []
+    client._spend_log_transactions_lock = asyncio.Lock()
+    client.jsonify_object = lambda data: dict(data)
+    client.db.is_connected = MagicMock(return_value=False)
+    client.db.connect = AsyncMock()
+    client.db.disconnect = AsyncMock()
+    client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}])
+    client.db.execute_raw = AsyncMock()
+    client.db.tx = MagicMock()
+    client.db.batch_ = MagicMock()
+    for table_name in _PRISMA_TABLES:
+        setattr(client.db, table_name, _make_table_mock())
+    return client
+
+
+@pytest.fixture
+def mock_dual_cache() -> MagicMock:
+    """In-memory DualCache stand-in.
+
+    Sync and async get/set wired against a private dict. Override or read
+    ``cache._store`` directly in a test for assertion convenience.
+    """
+    cache = MagicMock(name="MockDualCache")
+    cache._store: Dict[str, Any] = {}
+
+    def _sync_get(key: str, **_: Any) -> Any:
+        return cache._store.get(key)
+
+    def _sync_set(key: str, value: Any, **_: Any) -> None:
+        cache._store[key] = value
+
+    async def _async_get(key: str, **_: Any) -> Any:
+        return cache._store.get(key)
+
+    async def _async_set(key: str, value: Any, **_: Any) -> None:
+        cache._store[key] = value
+
+    async def _async_delete(key: str, **_: Any) -> None:
+        cache._store.pop(key, None)
+
+    cache.get_cache = MagicMock(side_effect=_sync_get)
+    cache.set_cache = MagicMock(side_effect=_sync_set)
+    cache.async_get_cache = AsyncMock(side_effect=_async_get)
+    cache.async_set_cache = AsyncMock(side_effect=_async_set)
+    cache.async_delete_cache = AsyncMock(side_effect=_async_delete)
+    return cache
+
+
+@pytest.fixture
+def patched_prisma_import(monkeypatch: pytest.MonkeyPatch) -> Iterator[MagicMock]:
+    """Replace ``prisma.Prisma`` and ``PrismaWrapper`` so PrismaClient.__init__
+    runs without a generated client. Yields the fake Prisma instance.
+
+    ``prisma`` raises RuntimeError (not AttributeError) for the missing
+    ``Prisma`` attribute, so ``monkeypatch.setattr`` can't probe it; assign
+    directly and restore in teardown.
+    """
+    import prisma as _prisma_pkg
+    import litellm.proxy.utils as _utils_mod
+
+    fake_prisma = MagicMock(name="FakePrisma")
+    fake_prisma.is_connected = MagicMock(return_value=False)
+    fake_prisma.connect = AsyncMock()
+    fake_prisma.disconnect = AsyncMock()
+
+    fake_prisma_factory = MagicMock(name="FakePrismaFactory", return_value=fake_prisma)
+    had_prisma_attr = "Prisma" in _prisma_pkg.__dict__
+    previous_prisma_attr = _prisma_pkg.__dict__.get("Prisma")
+    _prisma_pkg.Prisma = fake_prisma_factory  # type: ignore[attr-defined]
+
+    fake_wrapper = MagicMock(name="FakePrismaWrapper")
+    fake_wrapper.is_connected = MagicMock(return_value=False)
+    fake_wrapper.connect = AsyncMock()
+    fake_wrapper.disconnect = AsyncMock()
+    fake_wrapper.query_raw = AsyncMock(return_value=[{"?column?": 1}])
+
+    def _fake_wrapper_ctor(*args: Any, **kwargs: Any) -> MagicMock:
+        return fake_wrapper
+
+    monkeypatch.setattr(_utils_mod, "PrismaWrapper", _fake_wrapper_ctor)
+    fake_prisma.__wrapper__ = fake_wrapper
+    try:
+        yield fake_prisma
+    finally:
+        if had_prisma_attr:
+            _prisma_pkg.Prisma = previous_prisma_attr  # type: ignore[attr-defined]
+        else:
+            try:
+                del _prisma_pkg.Prisma  # type: ignore[attr-defined]
+            except AttributeError:
+                pass
+
+
+@pytest.fixture
+def prisma_client(
+    patched_prisma_import: MagicMock,
+    mock_prisma_client: MagicMock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> Any:
+    """Wired ``PrismaClient`` whose ``db`` attribute is the table mock.
+
+    The init runs through the real code path (testing the constructor's
+    config-attribute setup) and is then snapped to the easier-to-assert
+    table mock for downstream behavior pinning.
+    """
+    monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
+    monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False)
+    from litellm.proxy.utils import PrismaClient
+
+    proxy_logging_obj = MagicMock(name="MockProxyLogging")
+    proxy_logging_obj.failure_handler = AsyncMock()
+    pc = PrismaClient(
+        database_url="postgresql://test:test@localhost:5432/test",
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    pc.db = mock_prisma_client.db
+    return pc
+
+
+@dataclass
+class FakeClock:
+    """Monotonic-time controller for the spend monitor loop.
+
+    Tests advance time via ``clock.advance(seconds)`` while asyncio.sleep
+    is replaced with a clock-driven no-op.
+    """
+
+    now: float = 0.0
+    sleep_calls: List[float] = field(default_factory=list)
+
+    def advance(self, seconds: float) -> None:
+        self.now += seconds
+
+    def time(self) -> float:
+        return self.now
+
+    async def sleep(self, seconds: float) -> None:
+        self.sleep_calls.append(seconds)
+        self.now += seconds
+
+
+@pytest.fixture
+def fake_clock(monkeypatch: pytest.MonkeyPatch) -> FakeClock:
+    """Install a controllable clock + asyncio.sleep replacement."""
+    clock = FakeClock()
+    monkeypatch.setattr("time.time", clock.time)
+    monkeypatch.setattr("time.monotonic", clock.time)
+
+    async def _fast_sleep(seconds: float, *_: Any, **__: Any) -> None:
+        clock.sleep_calls.append(seconds)
+        clock.now += seconds
+
+    monkeypatch.setattr("asyncio.sleep", _fast_sleep)
+    return clock
+
+
+@pytest.fixture
+def make_spend_log_row() -> Callable[..., Dict[str, Any]]:
+    """Factory for fake LiteLLM_SpendLogs rows."""
+
+    def _make(
+        request_id: str = "req-1",
+        spend: float = 0.01,
+        model: str = "gpt-4o-mini",
+        **overrides: Any,
+    ) -> Dict[str, Any]:
+        row = {
+            "request_id": request_id,
+            "spend": spend,
+            "model": model,
+            "user": "user-1",
+            "team_id": "team-1",
+            "api_key": "hashed-key",
+            "startTime": "2026-06-02T00:00:00Z",
+            "endTime": "2026-06-02T00:00:01Z",
+            "metadata": {},
+        }
+        row.update(overrides)
+        return row
+
+    return _make
+
+
+@dataclass
+class _SentMessage:
+    from_addr: Optional[str]
+    to_addrs: Any
+    subject: Optional[str]
+    body: Optional[str]
+    starttls_called: bool
+    login_args: Optional[tuple]
+
+
+@dataclass
+class InMemorySMTP:
+    """Captures outbound SMTP traffic for ``send_email`` tests."""
+
+    sent: List[_SentMessage] = field(default_factory=list)
+    raise_on_send: Optional[Exception] = None
+
+    def server_factory(self) -> Callable[..., Any]:
+        outer = self
+
+        class _Conn:
+            def __init__(self) -> None:
+                self._starttls_called = False
+                self._login_args: Optional[tuple] = None
+
+            def __enter__(self) -> "_Conn":
+                return self
+
+            def __exit__(self, *exc: Any) -> None:
+                return None
+
+            def starttls(self) -> None:
+                self._starttls_called = True
+
+            def login(self, user: str, password: str) -> None:
+                self._login_args = (user, password)
+
+            def send_message(
+                self,
+                msg: EmailMessage,
+                from_addr: Optional[str] = None,
+                to_addrs: Any = None,
+            ) -> None:
+                if outer.raise_on_send is not None:
+                    raise outer.raise_on_send
+                body = ""
+                for part in msg.walk():
+                    if part.get_content_type() == "text/html":
+                        body = part.get_payload(decode=False) or ""
+                        break
+                outer.sent.append(
+                    _SentMessage(
+                        from_addr=from_addr,
+                        to_addrs=to_addrs,
+                        subject=msg["Subject"],
+                        body=body,
+                        starttls_called=self._starttls_called,
+                        login_args=self._login_args,
+                    )
+                )
+
+        def _factory(*args: Any, **kwargs: Any) -> _Conn:
+            return _Conn()
+
+        return _factory
+
+
+@pytest.fixture
+def in_memory_smtp(monkeypatch: pytest.MonkeyPatch) -> InMemorySMTP:
+    """Patch ``smtplib.SMTP`` to capture sends in memory.
+
+    Override ``smtp.raise_on_send`` to test the SMTP error path.
+    """
+    smtp = InMemorySMTP()
+    monkeypatch.setattr("smtplib.SMTP", smtp.server_factory())
+    return smtp

tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py

@@ -0,0 +1,400 @@
+"""Pin ``PrismaClient`` read-side data operations.
+
+Symbols pinned here:
+  - ``PrismaClient.hash_token``
+  - ``PrismaClient.jsonify_object``
+  - ``PrismaClient.jsonify_team_object``
+  - ``PrismaClient.check_view_exists``
+  - ``PrismaClient.get_request_status``
+  - ``PrismaClient.get_generic_data``
+  - ``PrismaClient._query_first_with_cached_plan_fallback``
+  - ``PrismaClient.get_data``
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import HTTPException
+
+from litellm.proxy.utils import PrismaClient
+
+
+def test_hash_token_method_returns_sha256(prisma_client: PrismaClient) -> None:
+    token = "sk-token-xyz"
+    actual = {
+        "result": prisma_client.hash_token(token),
+        "len": len(prisma_client.hash_token(token)),
+        "expected": hashlib.sha256(token.encode()).hexdigest(),
+        "deterministic": prisma_client.hash_token(token)
+        == prisma_client.hash_token(token),
+    }
+    assert actual == {
+        "result": hashlib.sha256(token.encode()).hexdigest(),
+        "len": 64,
+        "expected": hashlib.sha256(token.encode()).hexdigest(),
+        "deterministic": True,
+    }
+
+
+def test_hash_token_method_error_on_non_string(prisma_client: PrismaClient) -> None:
+    with pytest.raises(AttributeError):
+        prisma_client.hash_token(None)  # type: ignore[arg-type]
+
+
+def test_jsonify_object_serializes_nested_dicts(prisma_client: PrismaClient) -> None:
+    data = {
+        "metadata": {"a": 1, "b": [2, 3]},
+        "models": ["gpt-4o", "gpt-4o-mini"],
+        "token": "abc",
+        "spend": 1.23,
+    }
+    result = prisma_client.jsonify_object(data)
+    parsed_meta = json.loads(result["metadata"])
+    assert result == {
+        "metadata": json.dumps(data["metadata"]),
+        "models": ["gpt-4o", "gpt-4o-mini"],
+        "token": "abc",
+        "spend": 1.23,
+    }
+    assert parsed_meta == {"a": 1, "b": [2, 3]}
+
+
+def test_jsonify_object_fallback_for_unserializable_dict(
+    prisma_client: PrismaClient,
+) -> None:
+    class _Bad:
+        pass
+
+    data = {"metadata": {"x": _Bad()}, "label": "ok", "n": 1}
+    result = prisma_client.jsonify_object(data)
+    assert result == {
+        "metadata": "failed-to-serialize-json",
+        "label": "ok",
+        "n": 1,
+    }
+
+
+def test_jsonify_object_error_on_non_dict(prisma_client: PrismaClient) -> None:
+    with pytest.raises(AttributeError):
+        prisma_client.jsonify_object(None)  # type: ignore[arg-type]
+
+
+def test_jsonify_team_object_converts_members_to_json_string(
+    prisma_client: PrismaClient,
+) -> None:
+    data = {
+        "team_id": "t1",
+        "members_with_roles": [{"role": "admin", "user_id": "u1"}],
+        "metadata": {"foo": "bar"},
+        "models": ["gpt-4"],
+    }
+    result = prisma_client.jsonify_team_object(data)
+    assert result == {
+        "team_id": "t1",
+        "members_with_roles": json.dumps(data["members_with_roles"]),
+        "metadata": json.dumps(data["metadata"]),
+        "models": ["gpt-4"],
+    }
+
+
+def test_jsonify_team_object_error_on_non_dict(prisma_client: PrismaClient) -> None:
+    with pytest.raises(AttributeError):
+        prisma_client.jsonify_team_object(None)  # type: ignore[arg-type]
+
+
+@pytest.mark.parametrize(
+    "metadata,expected",
+    [
+        ({"status": "failure"}, "failure"),
+        ({"status": "success"}, "success"),
+        ({}, "success"),
+        ("not-json", "success"),
+        (json.dumps({"status": "failure"}), "failure"),
+    ],
+)
+def test_get_request_status_pins_status_resolution(
+    prisma_client: PrismaClient, metadata: Any, expected: str
+) -> None:
+    assert prisma_client.get_request_status({"metadata": metadata}) == expected
+
+
+def test_get_request_status_error_returns_success_default(
+    prisma_client: PrismaClient,
+) -> None:
+    """``get_request_status`` swallows AttributeError / JSONDecodeError and
+    defaults to ``success`` to avoid blocking the request pipeline.
+    """
+
+    class _Broken:
+        def get(self, *_: Any, **__: Any) -> Any:
+            raise AttributeError("broken metadata")
+
+    actual = prisma_client.get_request_status({"metadata": _Broken()})
+    assert actual == "success"
+
+
+@pytest.mark.asyncio
+async def test_get_generic_data_dispatches_by_table(
+    prisma_client: PrismaClient,
+) -> None:
+    row = SimpleNamespace(user_id="u1", spend=0.5, name="Alice")
+    prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=row)
+    result = await prisma_client.get_generic_data(
+        key="user_id", value="u1", table_name="users"
+    )
+    actual = {
+        "result_is_row": result is row,
+        "find_first_count": prisma_client.db.litellm_usertable.find_first.await_count,
+        "where_kwarg": prisma_client.db.litellm_usertable.find_first.await_args.kwargs[
+            "where"
+        ],
+        "user_attr": result.user_id,
+    }
+    assert actual == {
+        "result_is_row": True,
+        "find_first_count": 1,
+        "where_kwarg": {"user_id": "u1"},
+        "user_attr": "u1",
+    }
+
+
+@pytest.mark.asyncio
+async def test_get_generic_data_unknown_table_returns_none(
+    prisma_client: PrismaClient,
+) -> None:
+    result = await prisma_client.get_generic_data(
+        key="x", value="y", table_name="bogus"  # type: ignore[arg-type]
+    )
+    assert result is None
+
+
+@pytest.mark.asyncio
+async def test_get_generic_data_logs_failure_handler_and_raises_on_error(
+    prisma_client: PrismaClient,
+) -> None:
+    prisma_client.db.litellm_usertable.find_first = AsyncMock(
+        side_effect=RuntimeError("db boom")
+    )
+    with pytest.raises(RuntimeError, match="db boom"):
+        await prisma_client.get_generic_data(
+            key="user_id", value="x", table_name="users"
+        )
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_happy_returns_row(
+    prisma_client: PrismaClient,
+) -> None:
+    expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0}
+    prisma_client.db.query_first = AsyncMock(return_value=expected)
+    result = await prisma_client._query_first_with_cached_plan_fallback(
+        "SELECT * FROM x WHERE token = $1", "abc"
+    )
+    actual = {
+        "result": result,
+        "call_count": prisma_client.db.query_first.await_count,
+        "args": prisma_client.db.query_first.await_args.args,
+        "matches": result == expected,
+    }
+    assert actual == {
+        "result": expected,
+        "call_count": 1,
+        "args": ("SELECT * FROM x WHERE token = $1", "abc"),
+        "matches": True,
+    }
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_retries_on_cached_plan_error(
+    prisma_client: PrismaClient,
+) -> None:
+    expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0}
+    prisma_client.db.query_first = AsyncMock(
+        side_effect=[
+            RuntimeError("cached plan must not change result type"),
+            expected,
+        ]
+    )
+    result = await prisma_client._query_first_with_cached_plan_fallback(
+        "SELECT * FROM x WHERE token = $1", "abc"
+    )
+    assert result == expected
+    assert prisma_client.db.query_first.await_count == 2
+    second_call_sql = prisma_client.db.query_first.await_args_list[1].args[0]
+    assert "cache_invalidated_" in second_call_sql
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_reraises_non_plan_errors(
+    prisma_client: PrismaClient,
+) -> None:
+    prisma_client.db.query_first = AsyncMock(side_effect=RuntimeError("totally unrelated"))
+    with pytest.raises(RuntimeError, match="totally unrelated"):
+        await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
+
+
+@pytest.mark.asyncio
+async def test_check_view_exists_noop_when_all_views_present(
+    prisma_client: PrismaClient,
+) -> None:
+    prisma_client.db.query_raw = AsyncMock(
+        return_value=[
+            {
+                "view_count": 8,
+                "view_names": [
+                    "LiteLLM_VerificationTokenView",
+                    "MonthlyGlobalSpend",
+                    "Last30dKeysBySpend",
+                    "Last30dModelsBySpend",
+                    "MonthlyGlobalSpendPerKey",
+                    "MonthlyGlobalSpendPerUserPerKey",
+                    "Last30dTopEndUsersSpend",
+                    "DailyTagSpend",
+                ],
+            }
+        ]
+    )
+    prisma_client.db.execute_raw = AsyncMock()
+    result = await prisma_client.check_view_exists()
+    actual = {
+        "result": result,
+        "query_raw_calls": prisma_client.db.query_raw.await_count,
+        "execute_raw_calls": prisma_client.db.execute_raw.await_count,
+        "view_query_contains_token_view": "LiteLLM_VerificationTokenView"
+        in prisma_client.db.query_raw.await_args.args[0],
+    }
+    assert actual == {
+        "result": None,
+        "query_raw_calls": 1,
+        "execute_raw_calls": 0,
+        "view_query_contains_token_view": True,
+    }
+
+
+@pytest.mark.asyncio
+async def test_check_view_exists_creates_token_view_when_missing(
+    prisma_client: PrismaClient,
+) -> None:
+    prisma_client.db.query_raw = AsyncMock(
+        return_value=[
+            {
+                "view_count": 1,
+                "view_names": ["DailyTagSpend"],
+            }
+        ]
+    )
+    prisma_client.db.execute_raw = AsyncMock()
+    prisma_client.health_check = AsyncMock(return_value=[{"?column?": 1}])
+    result = await prisma_client.check_view_exists()
+    actual = {
+        "result": result,
+        "create_called": prisma_client.db.execute_raw.await_count,
+        "create_sql_starts_with_create_view": prisma_client.db.execute_raw.await_args.args[
+            0
+        ]
+        .strip()
+        .startswith('CREATE VIEW "LiteLLM_VerificationTokenView"'),
+    }
+    assert actual == {
+        "result": None,
+        "create_called": 1,
+        "create_sql_starts_with_create_view": True,
+    }
+
+
+@pytest.mark.asyncio
+async def test_check_view_exists_raises_when_query_raw_fails(
+    prisma_client: PrismaClient,
+) -> None:
+    prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("db down"))
+    with pytest.raises(RuntimeError, match="db down"):
+        await prisma_client.check_view_exists()
+
+
+@pytest.mark.asyncio
+async def test_get_data_token_find_unique_returns_record(
+    prisma_client: PrismaClient,
+) -> None:
+    token = "sk-key-1"
+    hashed = hashlib.sha256(token.encode()).hexdigest()
+    record = SimpleNamespace(token=hashed, user_id="u1", expires=None, spend=0.5)
+    prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        return_value=record
+    )
+
+    result = await prisma_client.get_data(token=token, table_name="key")
+    actual = {
+        "result_is_record": result is record,
+        "where_arg": prisma_client.db.litellm_verificationtoken.find_unique.await_args.kwargs[
+            "where"
+        ],
+        "include_arg": prisma_client.db.litellm_verificationtoken.find_unique.await_args.kwargs[
+            "include"
+        ],
+        "token_field_matches": result.token == hashed,
+    }
+    assert actual == {
+        "result_is_record": True,
+        "where_arg": {"token": hashed},
+        "include_arg": {"litellm_budget_table": True},
+        "token_field_matches": True,
+    }
+
+
+@pytest.mark.asyncio
+async def test_get_data_token_find_unique_missing_token_raises_401(
+    prisma_client: PrismaClient,
+) -> None:
+    prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
+    with pytest.raises(HTTPException) as excinfo:
+        await prisma_client.get_data(token="sk-missing", table_name="key")
+    err = excinfo.value
+    assert "invalid user key" in err.detail
+    assert err.status_code == 401
+
+
+@pytest.mark.asyncio
+async def test_get_data_user_find_unique_returns_user_row(
+    prisma_client: PrismaClient,
+) -> None:
+    row = SimpleNamespace(
+        user_id="u-7",
+        spend=1.5,
+        max_budget=10.0,
+        organization_memberships=[],
+    )
+    prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=row)
+    result = await prisma_client.get_data(user_id="u-7", table_name="user")
+    actual = {
+        "result_is_row": result is row,
+        "where_arg": prisma_client.db.litellm_usertable.find_unique.await_args.kwargs[
+            "where"
+        ],
+        "include_arg": prisma_client.db.litellm_usertable.find_unique.await_args.kwargs[
+            "include"
+        ],
+        "spend": row.spend,
+    }
+    assert actual == {
+        "result_is_row": True,
+        "where_arg": {"user_id": "u-7"},
+        "include_arg": {"organization_memberships": True},
+        "spend": 1.5,
+    }
+
+
+@pytest.mark.asyncio
+async def test_get_data_logs_and_raises_on_db_error(
+    prisma_client: PrismaClient,
+) -> None:
+    prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        side_effect=RuntimeError("network split")
+    )
+    with pytest.raises(RuntimeError, match="network split"):
+        await prisma_client.get_data(token="sk-broken", table_name="key")