diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py
index 8fe2e25a99..3504065a10 100644
--- a/tests/image_gen_tests/test_image_edits.py
+++ b/tests/image_gen_tests/test_image_edits.py
@@ -158,7 +158,7 @@ class TestAzureAIFlux2ImageEdit(BaseLLMImageEditTest):
         return {
             "model": "azure_ai/flux.2-pro",
             "image": SINGLE_TEST_IMAGE,
-            "api_base": "https://litellm-ci-cd-prod.services.ai.azure.com",
+            "api_base": os.getenv("AZURE_AI_API_BASE"),
             "api_key": os.getenv("AZURE_AI_API_KEY"),
             "api_version": "preview",
         }