Fix test case issue

[hsuyuming]

Title

Fix test case issue

Relevant issues

https://app.circleci.com/pipelines/github/BerriAI/litellm/32750/workflows/ac7ef503-4934-4b58-a55b-b6872e01838b/jobs/419628/tests#failed-test-0

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• [✓] I have Added testing in the tests/litellm/ directory, Adding at least 1 test is a hard requirement - see details
• [✓] I have added a screenshot of my new test passing locally
• [✓] My PR passes all unit tests on make test-unit
• [✓] My PR's scope is as isolated as possible, it only solves 1 specific problem

Type

🐛 Bug Fix
✅ Test

Changes

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
litellm	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	May 23, 2025 0:37am

[hsuyuming]

poetry run pytest tests/litellm/
=============================================== test session starts ===============================================
platform linux -- Python 3.10.12, pytest-7.4.4, pluggy-1.5.0
rootdir: /home/user/abehsu/litellm_clone
plugins: anyio-4.5.2, asyncio-0.21.2, mock-3.14.0, requests-mock-1.12.1, respx-0.22.0
asyncio: mode=strict
collected 874 items / 3 errors

===================================================== ERRORS ======================================================
_____________________ ERROR collecting tests/litellm/litellm_core_utils/test_token_counter.py _____________________
ImportError while importing test module '/home/user/abehsu/litellm_clone/tests/litellm/litellm_core_utils/test_token_counter.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
/usr/lib/python3.10/importlib/init.py:126: in import_module
return _bootstrap._gcd_import(name[level:], package, level)
tests/litellm/litellm_core_utils/test_token_counter.py:26: in
from tests.large_text import text
E ModuleNotFoundError: No module named 'tests.large_text'
__________________ ERROR collecting tests/litellm/litellm_core_utils/test_token_counter_tool.py ___________________
ImportError while importing test module '/home/user/abehsu/litellm_clone/tests/litellm/litellm_core_utils/test_token_counter_tool.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
/usr/lib/python3.10/importlib/init.py:126: in import_module
return _bootstrap._gcd_import(name[level:], package, level)
tests/litellm/litellm_core_utils/test_token_counter_tool.py:13: in
from test_token_counter import token_counter
tests/litellm/litellm_core_utils/test_token_counter.py:26: in
from tests.large_text import text
E ModuleNotFoundError: No module named 'tests.large_text'
_________________________ ERROR collecting tests/litellm/proxy/test_spend_log_cleanup.py __________________________
ImportError while importing test module '/home/user/abehsu/litellm_clone/tests/litellm/proxy/test_spend_log_cleanup.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
/usr/lib/python3.10/importlib/init.py:126: in import_module
return _bootstrap._gcd_import(name[level:], package, level)
tests/litellm/proxy/test_spend_log_cleanup.py:6: in
from datetime import datetime, timedelta, UTC, timezone
E ImportError: cannot import name 'UTC' from 'datetime' (/usr/lib/python3.10/datetime.py)
================================================ warnings summary =================================================
tests/litellm/integrations/test_custom_prompt_management.py:27
/home/user/abehsu/litellm_clone/tests/litellm/integrations/test_custom_prompt_management.py:27: PytestCollectionWarning: cannot collect test class 'TestCustomPromptManagement' because it has a init constructor (from: tests/litellm/integrations/test_custom_prompt_management.py)
class TestCustomPromptManagement(CustomPromptManagement):

tests/litellm/integrations/arize/test_arize_utils.py:181
/home/user/abehsu/litellm_clone/tests/litellm/integrations/arize/test_arize_utils.py:181: PytestCollectionWarning: cannot collect test class 'TestArizeLogger' because it has a init constructor (from: tests/litellm/integrations/arize/test_arize_utils.py)
class TestArizeLogger(CustomLogger):

tests/litellm/litellm_core_utils/test_streaming_handler.py:505
/home/user/abehsu/litellm_clone/tests/litellm/litellm_core_utils/test_streaming_handler.py:505: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(reruns=3)

tests/litellm/llms/ollama/test_ollama_chat_transformation.py:16
/home/user/abehsu/litellm_clone/tests/litellm/llms/ollama/test_ollama_chat_transformation.py:16: PytestCollectionWarning: cannot collect test class 'TestEvent' because it has a init constructor (from: tests/litellm/llms/ollama/test_ollama_chat_transformation.py)
class TestEvent(BaseModel):

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================= short test summary info =============================================
ERROR tests/litellm/litellm_core_utils/test_token_counter.py
ERROR tests/litellm/litellm_core_utils/test_token_counter_tool.py
ERROR tests/litellm/proxy/test_spend_log_cleanup.py
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Interrupted: 3 errors during collection !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
========================================== 4 warnings, 3 errors in 4.87s ==========================================
make: *** [Makefile:29: test-unit] Error 2

[krrishdholakia]

cc: @ishaan-jaff

[hsuyuming]

Hi @ishaan-jaff:
Can you help me review this PR when you get time? Thank you

[hsuyuming]

Hi @krrishdholakia @ishaan-jaff
We would like to having this feature soo if possible, can you help me to see whether this PR can be accepted and be be released to main?