diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py
index 08b5d1f93..fe5f8fa56 100644
--- a/litellm/tests/test_key_generate_prisma.py
+++ b/litellm/tests/test_key_generate_prisma.py
@@ -825,6 +825,7 @@ def test_call_with_key_over_budget_stream(prisma_client):
             )
             await track_cost_callback(
                 kwargs={
+                    "model": "sagemaker-chatgpt-v-2",
                     "stream": True,
                     "complete_streaming_response": resp,
                     "litellm_params": {
@@ -850,6 +851,7 @@ def test_call_with_key_over_budget_stream(prisma_client):
 
             assert spend_log.request_id == request_id
             assert spend_log.spend == float("5e-05")
+            assert spend_log.model == "sagemaker-chatgpt-v-2"
 
             # use generated key to auth in
             result = await user_api_key_auth(request=request, api_key=bearer_token)