From 23cad652569bfd6bb78ca8e9950ef0faeb8717ae Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 3 Aug 2024 20:28:11 -0700 Subject: [PATCH] test: fix test --- litellm/tests/test_image_generation.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_image_generation.py b/litellm/tests/test_image_generation.py index e59cf2865..fcbaa7457 100644 --- a/litellm/tests/test_image_generation.py +++ b/litellm/tests/test_image_generation.py @@ -194,6 +194,9 @@ async def test_aimage_generation_bedrock_with_optional_params(): pytest.fail(f"An exception occurred - {str(e)}") +from openai.types.image import Image + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_aimage_generation_vertex_ai(sync_mode): @@ -218,7 +221,7 @@ async def test_aimage_generation_vertex_ai(sync_mode): assert len(response.data) > 0 for d in response.data: - assert isinstance(d, litellm.ImageObject) + assert isinstance(d, Image) print("data in response.data", d) assert d.b64_json is not None except litellm.ServiceUnavailableError as e: