Update image tests

Author: moonbox3

python/tests/unit/connectors/ai/open_ai/services/test_azure_text_to_image.py

@@ -78,10 +78,12 @@ def test_azure_text_to_image_init_with_from_dict(azure_openai_unit_test_env) ->
         assert azure_text_to_image.client.default_headers[key] == value
 
 
-@patch.object(AsyncImages, "generate", return_value=AsyncMock(spec=ImagesResponse))
+@patch.object(AsyncImages, "generate")
 async def test_azure_text_to_image_calls_with_parameters(mock_generate, azure_openai_unit_test_env) -> None:
-    mock_generate.return_value.data = [Image(url="abc")]
-    mock_generate.return_value.usage = None
+    mock_response = AsyncMock(spec=ImagesResponse)
+    mock_response.data = [Image(url="abc")]
+    mock_response.usage = None
+    mock_generate.return_value = mock_response
 
     prompt = "A painting of a vase with flowers"
     width = 512

python/tests/unit/connectors/ai/open_ai/services/test_openai_text_to_image.py

@@ -32,10 +32,11 @@ def test_init(openai_unit_test_env):
     assert openai_text_to_image.ai_model_id == openai_unit_test_env["OPENAI_TEXT_TO_IMAGE_MODEL_ID"]
 
 
-def test_init_validation_fail() -> None:
+@pytest.mark.parametrize("exclude_list", [["OPENAI_TEXT_TO_IMAGE_MODEL_ID"]], indirect=True)
+def test_init_validation_fail(openai_unit_test_env) -> None:
     """Test that initialization fails when required parameters are missing."""
     with pytest.raises(ServiceInitializationError):
-        OpenAITextToImage(api_key="34523", ai_model_id=None)
+        OpenAITextToImage(api_key="34523", ai_model_id=None, env_file_path="test.env")
 
 
 def test_init_to_from_dict(openai_unit_test_env):
@@ -77,11 +78,13 @@ def test_prompt_execution_settings_class(openai_unit_test_env) -> None:
     assert openai_text_to_image.get_prompt_execution_settings_class() == OpenAITextToImageExecutionSettings
 
 
-@patch.object(AsyncImages, "generate", return_value=AsyncMock(spec=ImagesResponse))
+@patch.object(AsyncImages, "generate")
 async def test_generate_calls_with_parameters(mock_generate, openai_unit_test_env) -> None:
     """Test that generate_image calls the OpenAI API with correct parameters."""
-    mock_generate.return_value.data = [Image(url="abc")]
-    mock_generate.return_value.usage = None
+    mock_response = AsyncMock(spec=ImagesResponse)
+    mock_response.data = [Image(url="abc")]
+    mock_response.usage = None
+    mock_generate.return_value = mock_response
 
     ai_model_id = "test_model_id"
     prompt = "painting of flowers in vase"