test: fix flaky test

BerriAI · Jan 21, 2025 · 120308b · 120308b
1 parent 3086804
commit 120308b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/litellm/llms/huggingface/chat/handler.py b/litellm/llms/huggingface/chat/handler.py
@@ -432,6 +432,7 @@ def _transform_input(
         embed_url: str,
     ) -> dict:
         data: Dict = {}
+
         ## TRANSFORMATION ##
         if "sentence-transformers" in model:
             if len(input) == 0:

diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py
@@ -642,12 +642,16 @@ def tgi_mock_post(*args, **kwargs):
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
 
 
-@pytest.mark.parametrize("sync_mode", [True, False])
 @pytest.mark.asyncio
-async def test_hf_embedding_sentence_sim(sync_mode):
+@patch("litellm.llms.huggingface.chat.handler.async_get_hf_task_embedding_for_model")
+@pytest.mark.parametrize("sync_mode", [True, False])
+async def test_hf_embedding_sentence_sim(
+    mock_get_hf_task_embedding_for_model, sync_mode  # Add this parameter
+):
     try:
         # huggingface/microsoft/codebert-base
         # huggingface/facebook/bart-large
+        mock_get_hf_task_embedding_for_model.return_value = "sentence-similarity"
         if sync_mode is True:
             client = HTTPHandler(concurrent_limit=1)
         else: