diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index c967577be3d..ac18a65f287 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -151,6 +151,7 @@ async def test_get_response(): @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_get_router_response(): model = "claude-3-sonnet@20240229" vertex_ai_project = "adroit-crow-413218" @@ -195,6 +196,7 @@ async def test_get_router_response(): # @pytest.mark.skip( # reason="Local test. Vertex AI Quota is low. Leads to rate limit errors on ci/cd." # ) +@pytest.mark.flaky(retries=3, delay=1) def test_vertex_ai_anthropic(): model = "claude-3-sonnet@20240229" @@ -217,6 +219,7 @@ def test_vertex_ai_anthropic(): # @pytest.mark.skip( # reason="Local test. Vertex AI Quota is low. Leads to rate limit errors on ci/cd." # ) +@pytest.mark.flaky(retries=3, delay=1) def test_vertex_ai_anthropic_streaming(): try: load_vertex_ai_credentials() @@ -257,6 +260,7 @@ def test_vertex_ai_anthropic_streaming(): # reason="Local test. Vertex AI Quota is low. Leads to rate limit errors on ci/cd." # ) @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_vertex_ai_anthropic_async(): # load_vertex_ai_credentials() try: @@ -290,6 +294,7 @@ async def test_vertex_ai_anthropic_async(): # reason="Local test. Vertex AI Quota is low. Leads to rate limit errors on ci/cd." # ) @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_vertex_ai_anthropic_async_streaming(): # load_vertex_ai_credentials() try: @@ -324,6 +329,7 @@ async def test_vertex_ai_anthropic_async_streaming(): # asyncio.run(test_vertex_ai_anthropic_async_streaming()) +@pytest.mark.flaky(retries=3, delay=1) def test_vertex_ai(): import random @@ -374,6 +380,7 @@ def test_vertex_ai(): # test_vertex_ai() +@pytest.mark.flaky(retries=3, delay=1) def test_vertex_ai_stream(): load_vertex_ai_credentials() litellm.set_verbose = True @@ -422,6 +429,7 @@ def test_vertex_ai_stream(): # test_vertex_ai_stream() +@pytest.mark.flaky(retries=3, delay=1) @pytest.mark.asyncio async def test_async_vertexai_response(): import random @@ -464,6 +472,7 @@ async def test_async_vertexai_response(): # asyncio.run(test_async_vertexai_response()) +@pytest.mark.flaky(retries=3, delay=1) @pytest.mark.asyncio async def test_async_vertexai_streaming_response(): import random @@ -519,6 +528,7 @@ async def test_async_vertexai_streaming_response(): @pytest.mark.parametrize("provider", ["vertex_ai"]) # "vertex_ai_beta" @pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.flaky(retries=3, delay=1) @pytest.mark.asyncio async def test_gemini_pro_vision(provider, sync_mode): try: @@ -584,6 +594,7 @@ async def test_gemini_pro_vision(provider, sync_mode): @pytest.mark.parametrize("load_pdf", [False]) # True, +@pytest.mark.flaky(retries=3, delay=1) def test_completion_function_plus_pdf(load_pdf): litellm.set_verbose = True load_vertex_ai_credentials() @@ -926,6 +937,7 @@ from litellm.tests.test_completion import response_format_tests "sync_mode", [True, False], ) # +@pytest.mark.flaky(retries=3, delay=1) @pytest.mark.asyncio async def test_partner_models_httpx(model, sync_mode): try: @@ -981,6 +993,7 @@ async def test_partner_models_httpx(model, sync_mode): [True, False], # ) # @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_partner_models_httpx_streaming(model, sync_mode): try: load_vertex_ai_credentials() @@ -1157,6 +1170,7 @@ def vertex_httpx_mock_post(url, data=None, json=None, headers=None): @pytest.mark.parametrize("provider", ["vertex_ai_beta"]) # "vertex_ai", @pytest.mark.parametrize("content_filter_type", ["prompt", "response"]) # "vertex_ai", @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_gemini_pro_json_schema_httpx_content_policy_error( provider, content_filter_type ): @@ -1629,6 +1643,7 @@ async def test_gemini_pro_httpx_custom_api_base(provider): @pytest.mark.parametrize("sync_mode", [True]) @pytest.mark.parametrize("provider", ["vertex_ai"]) @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_gemini_pro_function_calling(provider, sync_mode): try: load_vertex_ai_credentials() @@ -1714,6 +1729,7 @@ async def test_gemini_pro_function_calling(provider, sync_mode): @pytest.mark.parametrize("sync_mode", [True]) @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_gemini_pro_function_calling_streaming(sync_mode): load_vertex_ai_credentials() litellm.set_verbose = True @@ -1779,6 +1795,7 @@ async def test_gemini_pro_function_calling_streaming(sync_mode): @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_gemini_pro_async_function_calling(): load_vertex_ai_credentials() try: @@ -1831,6 +1848,7 @@ async def test_gemini_pro_async_function_calling(): # asyncio.run(gemini_pro_async_function_calling()) +@pytest.mark.flaky(retries=3, delay=1) def test_vertexai_embedding(): try: load_vertex_ai_credentials() @@ -1937,6 +1955,7 @@ def test_vertexai_embedding_embedding_latest(): @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_vertexai_aembedding(): try: load_vertex_ai_credentials()