From 67433a04a250e7e33f1cccbdedd95c74c0db35f1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 2 Jul 2024 22:12:44 -0700 Subject: [PATCH] test: fix test --- litellm/tests/test_router.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_router.py b/litellm/tests/test_router.py index 7c59611d730..43525afcb21 100644 --- a/litellm/tests/test_router.py +++ b/litellm/tests/test_router.py @@ -1865,7 +1865,10 @@ async def test_router_model_usage(mock_response): ) await asyncio.sleep(3) - initial_usage = await llm_router.get_model_group_usage(model_group=model) + initial_usage_tuple = await llm_router.get_model_group_usage( + model_group=model + ) + initial_usage = initial_usage_tuple[0] # completion call - 10 tokens _ = await llm_router.acompletion( @@ -1873,7 +1876,10 @@ async def test_router_model_usage(mock_response): ) await asyncio.sleep(3) - updated_usage = await llm_router.get_model_group_usage(model_group=model) + updated_usage_tuple = await llm_router.get_model_group_usage( + model_group=model + ) + updated_usage = updated_usage_tuple[0] assert updated_usage == initial_usage + 10 # type: ignore break