diff --git a/tests/test_litellm/test_dashscope_image_generation.py b/tests/test_litellm/test_dashscope_image_generation.py index 816b11fda1e..b7680a7e7fd 100644 --- a/tests/test_litellm/test_dashscope_image_generation.py +++ b/tests/test_litellm/test_dashscope_image_generation.py @@ -326,66 +326,3 @@ def test_litellm_image_generation_dashscope_end_to_end(): assert "input" in body assert "messages" in body["input"] - -# --------------------------------------------------------------------------- -# 7. Auto-redirect: litellm.acompletion() → aimage_generation() for image models -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_acompletion_auto_redirects_to_image_generation(): - """ - When litellm.acompletion() is called with a qwen-image model, - it should detect mode='image_generation' and auto-redirect to - aimage_generation(), returning a ModelResponse with Markdown image content. - """ - mock_response_body = { - "output": { - "choices": [ - { - "finish_reason": "stop", - "message": { - "role": "assistant", - "content": [{"image": "https://dashscope-result.oss.aliyuncs.com/auto.png"}], - }, - } - ] - }, - "usage": {"input_tokens": 0, "output_tokens": 0, "width": 1024, "height": 1024, "image_count": 1}, - } - - import os - - prev_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") - prev_model_cost = litellm.model_cost - try: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: - mock_http_response = MagicMock() - mock_http_response.json.return_value = mock_response_body - mock_http_response.status_code = 200 - mock_http_response.headers = {} - mock_post.return_value = mock_http_response - - response = await litellm.acompletion( - model="dashscope/qwen-image-2.0", - messages=[{"role": "user", "content": "a puppy on green grass"}], - api_key="sk-test-key", - ) - - assert response is not None - assert response.object == "chat.completion" - assert len(response.choices) == 1 - content = response.choices[0].message.content - assert "https://dashscope-result.oss.aliyuncs.com/auto.png" in content - assert "![image](" in content - finally: - if prev_env is None: - os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = prev_env - litellm.model_cost = prev_model_cost