From a356b2f81bda8418b4d2fc3e1a7bdfc251d5a8e3 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 6 May 2026 16:05:04 -0700
Subject: [PATCH] test(agent_session): RecordingNoopProvider tracks
 provision/terminate calls; warm-pool tables in fake DB

---
 .../proxy/agent_session_endpoints/conftest.py | 48 +++++++++++++++++--
 1 file changed, 45 insertions(+), 3 deletions(-)

diff --git a/tests/test_litellm/proxy/agent_session_endpoints/conftest.py b/tests/test_litellm/proxy/agent_session_endpoints/conftest.py
index 7c76336608e..09536f68e3d 100644
--- a/tests/test_litellm/proxy/agent_session_endpoints/conftest.py
+++ b/tests/test_litellm/proxy/agent_session_endpoints/conftest.py
@@ -182,6 +182,11 @@ class _DB:
         self.litellm_agentsession = _Table()
         self.litellm_agentrun = _AgentRunTable()
         self.litellm_agentrunevent = _AgentRunEventTable()
+        # Warm-pool tables (LIT-2890). Same `_Table` semantics so the in-memory
+        # client mirrors what the real Prisma client exposes.
+        self.litellm_agentvm = _Table()
+        self.litellm_agentvmconfig = _Table()
+        self.litellm_agentsecret = _Table()
 
 
 class _AgentRunTable(_Table):
@@ -290,11 +295,48 @@ def other_tenant_client(fake_prisma_client):
     return _build_test_app(LitellmUserRoles.INTERNAL_USER, api_key="sk-tenant-B")
 
 
+class _RecordingNoopProvider:
+    """Test-only noop that tracks ``provision`` / ``terminate`` calls.
+
+    A1's tests poke at ``provider.provision_calls`` / ``terminate_calls``
+    while B1's ``NoopProvider`` only implements the ABC. Wrap B1's behavior
+    here without polluting the production class.
+    """
+
+    name = "noop"
+
+    def __init__(self) -> None:
+        from litellm.proxy.agent_session_endpoints.vm_providers import NoopProvider
+
+        self._inner = NoopProvider()
+        self.provision_calls: list = []
+        self.terminate_calls: list = []
+
+    async def provision(self, ctx) -> Any:  # noqa: ANN001
+        self.provision_calls.append(
+            {
+                "session_id": getattr(ctx, "session_id", None),
+                "team_id": getattr(ctx, "team_id", None),
+                "agent_id": getattr(ctx, "agent_id", None),
+                "mode": getattr(ctx, "mode", None),
+            }
+        )
+        return await self._inner.provision(ctx)
+
+    async def terminate(self, vm) -> None:  # noqa: ANN001
+        self.terminate_calls.append(
+            {"session_id": vm.metadata.get("session_id"), "vm_id": vm.vm_id}
+        )
+        await self._inner.terminate(vm)
+
+    async def status(self, vm):  # noqa: ANN001
+        return await self._inner.status(vm)
+
+
 @pytest.fixture
 def noop_provider(monkeypatch):
-    """Reset the VM provider registry to a fresh ``NoopVMProvider``."""
+    """Reset the VM provider registry to a fresh recording noop."""
     from litellm.proxy.agent_session_endpoints.vm_providers import (
-        NoopVMProvider,
         register_vm_provider,
     )
     from litellm.proxy.agent_session_endpoints.vm_providers.registry import (
@@ -302,7 +344,7 @@ def noop_provider(monkeypatch):
     )
 
     reset_vm_provider_registry()
-    provider = NoopVMProvider()
+    provider = _RecordingNoopProvider()
     register_vm_provider(provider)
     yield provider
     reset_vm_provider_registry()