diff --git a/tests/unit/test_rate_limit_error_unification.py b/tests/unit/test_rate_limit_error_unification.py
index e5acba938c7..0e0e5273e6d 100644
--- a/tests/unit/test_rate_limit_error_unification.py
+++ b/tests/unit/test_rate_limit_error_unification.py
@@ -498,6 +498,7 @@ class TestProxyHooksActuallyRaiseProxyRateLimitError:
             "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
         ) as mock_registry:
             mock_registry.get_agent_by_id.return_value = agent
+            mock_registry.stable_agent_id.return_value = "agent-iter-1"
             # First call within budget.
             await handler.async_pre_call_hook(
                 user_api_key_dict=user_api_key_dict,
@@ -854,6 +855,7 @@ class TestProxyHooksActuallyRaiseProxyRateLimitError:
             "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
         ) as mock_registry:
             mock_registry.get_agent_by_id.return_value = agent
+            mock_registry.stable_agent_id.return_value = "agent-session-1"
             with pytest.raises(ProxyRateLimitError) as exc_info:
                 await handler.async_pre_call_hook(
                     user_api_key_dict=user_api_key_dict,