From 0141ea4a3191ce04ecab8663ed72b8d049138941 Mon Sep 17 00:00:00 2001
From: ryan <ryan@berri.ai>
Date: Sun, 20 Sep 2026 10:11:02 +0000
Subject: [PATCH] test(proxy): inject the loaded catalog into update_db_model
 instead of patching the class

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../model_management_endpoints.py             | 19 ++++++++++++++-----
 .../test_model_management_endpoints.py        |  6 ++----
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py
index f3acd794309..d76d533f996 100644
--- a/litellm/proxy/management_endpoints/model_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/model_management_endpoints.py
@@ -19,7 +19,7 @@ from dataclasses import dataclass
 from fnmatch import fnmatchcase
 from json import JSONDecodeError
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast, runtime_checkable
+from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable
 
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
 from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
@@ -888,14 +888,23 @@ def _cost_map_entry(db_model: Deployment, incoming_model_info: Mapping[str, obje
         return MappingProxyType({})
 
 
-def _loaded_catalog_entry(incoming_model_info: Mapping[str, object]) -> Mapping[str, object]:
+LoadedCatalog: TypeAlias = Callable[[], Mapping[str, Mapping[str, object]]]  # mutable-ok: Callable parameter syntax
+
+
+def _loaded_catalog_entry(
+    incoming_model_info: Mapping[str, object], loaded_catalog: LoadedCatalog
+) -> Mapping[str, object]:
     catalog_key: Final = incoming_model_info.get(COST_MAP_LOOKUP_KEY)
     if not isinstance(catalog_key, str):
         return MappingProxyType({})
-    return GetModelCostMap.loaded_model_cost_map().get(catalog_key, MappingProxyType({}))
+    return loaded_catalog().get(catalog_key, MappingProxyType({}))
 
 
-def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel:
+def update_db_model(
+    db_model: Deployment,
+    updated_patch: updateDeployment,
+    loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map,
+) -> PrismaCompatibleUpdateDBModel:
     if updated_patch.model_info is not None:
         _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True))
     merged_model_name: Final = updated_patch.model_name or db_model.model_name
@@ -921,7 +930,7 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
         echoed_fields: Final = echoed_cost_map_fields(
             incoming_model_info,
             _cost_map_entry(db_model, incoming_model_info),
-            _loaded_catalog_entry(incoming_model_info),
+            _loaded_catalog_entry(incoming_model_info, loaded_catalog),
         )
         merged_model_info.update(
             MappingProxyType(
diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
index cf37c99322e..cb7c960daf5 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
@@ -4119,15 +4119,12 @@ class TestModelInfoCostMapEchoFilter:
 
         import litellm
 
-        from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
         from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
         from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
 
         bundled = litellm.get_model_info("openai/gpt-5.6")
         remote = {**bundled, "max_input_tokens": bundled["max_input_tokens"] + 1}
-        monkeypatch.setattr(
-            GetModelCostMap, "_loaded_catalog", MappingProxyType({remote["key"]: MappingProxyType(remote)})
-        )
+        remote_catalog = MappingProxyType({remote["key"]: MappingProxyType(remote)})
         monkeypatch.setattr(litellm, "get_model_info", lambda model, **_: {**remote, "max_input_tokens": 2048})
         db_model = Deployment(
             model_name="gpt-5.6",
@@ -4138,6 +4135,7 @@ class TestModelInfoCostMapEchoFilter:
         result = update_db_model(
             db_model=db_model,
             updated_patch=updateDeployment(model_info=ModelInfo(**{**remote, "id": "dep-echo-9", "db_model": True})),
+            loaded_catalog=lambda: remote_catalog,
         )
 
         info = json.loads(result["model_info"])