From 68fe044b32e621c4efb8e68eed0fdb741fe7b2fc Mon Sep 17 00:00:00 2001
From: Mrinal Chanshetty <mrinal@berri.ai>
Date: Tue, 20 Oct 2026 16:09:19 -0700
Subject: [PATCH] test: add the connected_proxy_prisma integration helper

Prerequisite for the #47190 integration test on stable/1.106.x. Copied unchanged from tests/integration/_support/prisma_clients.py on main, added there by #45939.
---
 tests/integration/_support/prisma_clients.py | 35 ++++++++++++++++++++
 1 file changed, 35 insertions(+)
 create mode 100644 tests/integration/_support/prisma_clients.py

diff --git a/tests/integration/_support/prisma_clients.py b/tests/integration/_support/prisma_clients.py
new file mode 100644
index 00000000000..5f765e2a6db
--- /dev/null
+++ b/tests/integration/_support/prisma_clients.py
@@ -0,0 +1,35 @@
+import os
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Final
+
+from prisma import Prisma
+from prisma.types import DatasourceOverride
+
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.utils import PrismaClient, ProxyLogging
+
+
+@asynccontextmanager
+async def connected_prisma(database_url: str | None = None) -> AsyncIterator[Prisma]:
+    client: Final = Prisma(
+        datasource=DatasourceOverride(url=database_url)
+    ) if database_url is not None else Prisma()
+    await client.connect()
+    try:
+        yield client
+    finally:
+        await client.disconnect()
+
+
+@asynccontextmanager
+async def connected_proxy_prisma(database_url: str | None = None) -> AsyncIterator[PrismaClient]:
+    client: Final = PrismaClient(
+        database_url=database_url or os.environ["DATABASE_URL"],
+        proxy_logging_obj=ProxyLogging(UserApiKeyCache()),
+    )
+    await client.connect()
+    try:
+        yield client
+    finally:
+        await client.disconnect()