Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(image): honor provider proxy for URL downloads

Browse Source
This commit is contained in:
chengyongru
2026-07-27 10:06:19 +08:00
committed by Xubin Ren
parent cc3dbbe804
commit d73794bc68
4 changed files with 211 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/providers/test_image_generation.py
+44 -14
View File
@@ -103,19 +103,19 @@ class CodexStreamingCompleteThenErrorResponse(FakeResponse):
@pytest.fixture(autouse=True)
def generated_image_downloads(monkeypatch) -> list[str]:
def generated_image_downloads(monkeypatch) -> list[tuple[str, str | None]]:
"""Keep provider response parsing tests independent from outbound HTTP."""
urls: list[str] = []
downloads: list[tuple[str, str | None]] = []
async def download(url: str) -> str:
urls.append(url)
async def download(url: str, *, proxy: str | None = None) -> str:
downloads.append((url, proxy))
return PNG_DATA_URL
monkeypatch.setattr(
"nanobot.providers.image_generation._download_image_data_url",
download,
)
return urls
return downloads
@pytest.mark.asyncio
@@ -294,19 +294,21 @@ async def test_aihubmix_image_edit_payload_uses_reference_images(tmp_path: Path)
@pytest.mark.asyncio
async def test_aihubmix_image_generation_downloads_url_response(
generated_image_downloads: list[str],
generated_image_downloads: list[tuple[str, str | None]],
) -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
proxy = "http://127.0.0.1:23458"
client = AIHubMixImageGenerationClient(
api_key="sk-ahm-test",
proxy=proxy,
client=fake, # type: ignore[arg-type]
)
response = await client.generate(prompt="draw", model="gpt-image-2-free")
assert response.images[0].startswith("data:image/png;base64,")
assert generated_image_downloads == ["https://cdn.example/image.png"]
assert generated_image_downloads == [("https://cdn.example/image.png", proxy)]
@pytest.mark.asyncio
@@ -836,18 +838,22 @@ async def test_openai_b64_json_response_uses_detected_mime() -> None:
@pytest.mark.asyncio
async def test_openai_url_download_fallback(generated_image_downloads: list[str]) -> None:
async def test_openai_url_download_fallback(
generated_image_downloads: list[tuple[str, str | None]],
) -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
proxy = "http://127.0.0.1:23458"
client = OpenAIImageGenerationClient(
api_key="sk-openai-test",
proxy=proxy,
client=fake, # type: ignore[arg-type]
)
response = await client.generate(prompt="draw", model="dall-e-3")
assert response.images[0].startswith("data:image/png;base64,")
assert generated_image_downloads == ["https://cdn.example/image.png"]
assert generated_image_downloads == [("https://cdn.example/image.png", proxy)]
@pytest.mark.asyncio
@@ -1211,14 +1217,16 @@ async def test_custom_generate_maps_one_k_to_openai_dimension() -> None:
@pytest.mark.asyncio
async def test_custom_generate_extra_body_can_override_defaults(
generated_image_downloads: list[str],
generated_image_downloads: list[tuple[str, str | None]],
) -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://images.example/cat.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
proxy = "http://127.0.0.1:23458"
client = CustomImageGenerationClient(
api_key="sk-custom-test",
api_base="https://custom.example/v1",
extra_body={"response_format": "url", "size": "2K"},
proxy=proxy,
client=fake, # type: ignore[arg-type]
)
@@ -1229,7 +1237,7 @@ async def test_custom_generate_extra_body_can_override_defaults(
)
assert response.images == [PNG_DATA_URL]
assert generated_image_downloads == ["https://images.example/cat.png"]
assert generated_image_downloads == [("https://images.example/cat.png", proxy)]
body = fake.calls[0]["json"]
assert body["response_format"] == "url"
assert body["size"] == "2K"
@@ -1636,19 +1644,21 @@ async def test_zhipu_image_generation_with_explicit_size() -> None:
@pytest.mark.asyncio
async def test_zhipu_image_generation_downloads_url_response(
generated_image_downloads: list[str],
generated_image_downloads: list[tuple[str, str | None]],
) -> None:
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
fake.get_response = FakeResponse({}, content=PNG_BYTES)
proxy = "http://127.0.0.1:23458"
client = ZhipuImageGenerationClient(
api_key="sk-zhipu-test",
proxy=proxy,
client=fake, # type: ignore[arg-type]
)
response = await client.generate(prompt="draw", model="glm-image")
assert response.images[0].startswith("data:image/png;base64,")
assert generated_image_downloads == ["https://cdn.example/image.png"]
assert generated_image_downloads == [("https://cdn.example/image.png", proxy)]
@pytest.mark.asyncio
@@ -1728,7 +1738,9 @@ def _modelscope_fast_poll(monkeypatch) -> None:
@pytest.mark.asyncio
async def test_modelscope_image_generation_submit_and_poll() -> None:
async def test_modelscope_image_generation_submit_and_poll(
generated_image_downloads: list[tuple[str, str | None]],
) -> None:
submit = FakeResponse({"task_id": "abc123"})
poll_responses = [
FakeResponse({"task_status": "PENDING"}),
@@ -1738,9 +1750,11 @@ async def test_modelscope_image_generation_submit_and_poll() -> None:
}),
]
fake = ModelScopeFakeClient(submit, poll_responses)
proxy = "http://127.0.0.1:23458"
client = ModelScopeImageGenerationClient(
api_key="ms-token",
api_base="https://api-inference.modelscope.cn/v1",
proxy=proxy,
client=fake, # type: ignore[arg-type]
)
@@ -1750,6 +1764,7 @@ async def test_modelscope_image_generation_submit_and_poll() -> None:
)
assert response.images[0].startswith("data:image/png;base64,")
assert generated_image_downloads == [("https://cdn.example/image.png", proxy)]
# Verify POST request
post_call = fake.calls[0]
@@ -1919,3 +1934,18 @@ async def test_modelscope_image_generation_poll_timeout(monkeypatch) -> None:
# Should have polled up to the (patched) attempt limit.
assert len(fake.get_calls) == 3
def test_image_provider_http_client_kwargs_include_explicit_proxy() -> None:
proxy = "http://127.0.0.1:23458"
client = AIHubMixImageGenerationClient(
api_key="sk-ahm-test",
proxy=proxy,
)
assert client._http_client_kwargs() == {
"timeout": client.timeout,
"proxy": proxy,
"trust_env": False,
}
tests/providers/test_image_generation_security.py
+110 -1
View File
@@ -33,8 +33,16 @@ def _resolve_public(host: str, port: int | None, *args, **kwargs):
["http://127.0.0.1/admin", "http://[::]/admin"],
ids=["ipv4-loopback", "ipv6-unspecified"],
)
@pytest.mark.parametrize(
"proxy",
[None, "http://127.0.0.1:23458"],
ids=["direct", "explicit-proxy"],
)
@pytest.mark.asyncio
async def test_generated_image_download_blocks_unsafe_target(url: str) -> None:
async def test_generated_image_download_blocks_unsafe_target(
url: str,
proxy: str | None,
) -> None:
requested = False
async def handler(request: httpx.Request) -> httpx.Response:
@@ -45,6 +53,7 @@ async def test_generated_image_download_blocks_unsafe_target(url: str) -> None:
with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"):
await _download_image_data_url(
url,
proxy=proxy,
transport=httpx.MockTransport(handler),
)
@@ -116,3 +125,103 @@ async def test_generated_image_download_enforces_streaming_size_limit(monkeypatc
"https://cdn.example/image.png",
transport=httpx.MockTransport(handler),
)
class _StreamContext:
def __init__(self, response: httpx.Response) -> None:
self.response = response
async def __aenter__(self) -> httpx.Response:
return self.response
async def __aexit__(self, exc_type, exc, traceback) -> None:
await self.response.aclose()
@pytest.mark.asyncio
async def test_generated_image_download_uses_explicit_provider_proxy(
monkeypatch,
) -> None:
monkeypatch.setattr(
"nanobot.security.network.socket.getaddrinfo",
_resolve_public,
)
captured: dict[str, object] = {}
class FakeAsyncClient:
def __init__(self, **kwargs) -> None:
captured["kwargs"] = kwargs
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, traceback) -> None:
return None
def stream(self, method: str, url: str) -> _StreamContext:
captured["request"] = (method, url)
request = httpx.Request(method, url)
return _StreamContext(httpx.Response(200, content=PNG_BYTES, request=request))
monkeypatch.setattr(image_generation.httpx, "AsyncClient", FakeAsyncClient)
proxy = "http://127.0.0.1:23458"
result = await _download_image_data_url(
"https://cdn.example/image.png",
proxy=proxy,
)
assert result.startswith("data:image/png;base64,")
assert captured["request"] == ("GET", "https://cdn.example/image.png")
assert captured["kwargs"] == {
"follow_redirects": False,
"timeout": image_generation._DEFAULT_TIMEOUT_S,
"trust_env": False,
"proxy": proxy,
}
@pytest.mark.asyncio
async def test_proxied_generated_image_download_revalidates_redirects(
monkeypatch,
) -> None:
original_getaddrinfo = socket.getaddrinfo
def resolve_test_hosts(host: str, port: int | None, *args, **kwargs):
if host == "cdn.example":
return _resolve_public(host, port, *args, **kwargs)
return original_getaddrinfo(host, port, *args, **kwargs)
monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", resolve_test_hosts)
requested: list[str] = []
class FakeAsyncClient:
def __init__(self, **kwargs) -> None:
pass
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, traceback) -> None:
return None
def stream(self, method: str, url: str) -> _StreamContext:
requested.append(url)
request = httpx.Request(method, url)
return _StreamContext(
httpx.Response(
302,
headers={"location": "http://169.254.169.254/latest"},
request=request,
)
)
monkeypatch.setattr(image_generation.httpx, "AsyncClient", FakeAsyncClient)
with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"):
await _download_image_data_url(
"https://cdn.example/image.png",
proxy="http://127.0.0.1:23458",
)
assert requested == ["https://cdn.example/image.png"]