diff --git a/astrbot/core/utils/io.py b/astrbot/core/utils/io.py index 41cc87639a..a2f7946e8d 100644 --- a/astrbot/core/utils/io.py +++ b/astrbot/core/utils/io.py @@ -113,6 +113,20 @@ def save_temp_img(img: Image.Image | bytes) -> str: return p +async def _save_image_response( + resp: aiohttp.ClientResponse, + url: str, + path: str | None, +) -> str: + _raise_for_download_status(resp, url) + content = await resp.read() + if not path: + return save_temp_img(content) + with open(path, "wb") as f: + f.write(content) + return path + + async def download_image_by_url( url: str, post: bool = False, @@ -131,18 +145,10 @@ async def download_image_by_url( ) as session: if post: async with session.post(url, json=post_data) as resp: - if not path: - return save_temp_img(await resp.read()) - with open(path, "wb") as f: - f.write(await resp.read()) - return path + return await _save_image_response(resp, url, path) else: async with session.get(url) as resp: - if not path: - return save_temp_img(await resp.read()) - with open(path, "wb") as f: - f.write(await resp.read()) - return path + return await _save_image_response(resp, url, path) except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError): # 关闭SSL验证(仅在证书验证失败时作为fallback) logger.warning( @@ -157,20 +163,10 @@ async def download_image_by_url( async with aiohttp.ClientSession() as session: if post: async with session.post(url, json=post_data, ssl=ssl_context) as resp: - if not path: - return save_temp_img(await resp.read()) - with open(path, "wb") as f: - f.write(await resp.read()) - return path + return await _save_image_response(resp, url, path) else: async with session.get(url, ssl=ssl_context) as resp: - if not path: - return save_temp_img(await resp.read()) - with open(path, "wb") as f: - f.write(await resp.read()) - return path - except Exception as e: - raise e + return await _save_image_response(resp, url, path) async def _emit_download_progress(progress_callback, payload: dict) -> None: diff --git a/tests/unit/test_io_download_file.py b/tests/unit/test_io_download_file.py index 5cea2f0e50..71aa125e7b 100644 --- a/tests/unit/test_io_download_file.py +++ b/tests/unit/test_io_download_file.py @@ -17,8 +17,12 @@ class _FakeResponse: def __init__(self, *, status: int, chunks: list[bytes]): self.status = status self.headers = {"content-length": str(sum(len(chunk) for chunk in chunks))} + self._body = b"".join(chunks) self.content = _FakeContent(chunks) + async def read(self) -> bytes: + return self._body + async def __aenter__(self): return self @@ -41,6 +45,11 @@ def get(self, *_args, **_kwargs): raise self._response return self._response + def post(self, *_args, **_kwargs): + if isinstance(self._response, Exception): + raise self._response + return self._response + def _patch_download_session(monkeypatch, response: _FakeResponse): _patch_download_sessions(monkeypatch, [response]) @@ -105,3 +114,83 @@ async def test_download_file_writes_successful_response(monkeypatch, tmp_path): await io.download_file("https://example.test/ok.bin", str(target_path)) assert target_path.read_bytes() == b"hello world" + + +@pytest.mark.parametrize("post", [False, True]) +@pytest.mark.asyncio +async def test_download_image_by_url_rejects_non_200_response( + monkeypatch, + tmp_path, + post, +): + target_path = tmp_path / "image.jpg" + _patch_download_session( + monkeypatch, + _FakeResponse(status=404, chunks=[b'{"error":"not found"}']), + ) + + with pytest.raises(io.DownloadFileHTTPError, match="HTTP status code: 404"): + await io.download_image_by_url( + "https://example.test/generate", + post=post, + post_data={"json": False} if post else None, + path=str(target_path), + ) + + assert not target_path.exists() + + +@pytest.mark.parametrize("post", [False, True]) +@pytest.mark.asyncio +async def test_download_image_by_url_rejects_non_200_response_after_ssl_fallback( + monkeypatch, + tmp_path, + post, +): + class FakeSSLError(Exception): + pass + + target_path = tmp_path / "image.jpg" + _patch_download_sessions( + monkeypatch, + [ + FakeSSLError(), + _FakeResponse(status=500, chunks=[b"server error"]), + ], + ) + monkeypatch.setattr(io.aiohttp, "ClientConnectorSSLError", FakeSSLError) + monkeypatch.setattr(io.aiohttp, "ClientConnectorCertificateError", FakeSSLError) + + with pytest.raises(io.DownloadFileHTTPError, match="HTTP status code: 500"): + await io.download_image_by_url( + "https://example.test/generate", + post=post, + post_data={"json": False} if post else None, + path=str(target_path), + ) + + assert not target_path.exists() + + +@pytest.mark.parametrize("post", [False, True]) +@pytest.mark.asyncio +async def test_download_image_by_url_writes_successful_response( + monkeypatch, + tmp_path, + post, +): + target_path = tmp_path / "image.jpg" + _patch_download_session( + monkeypatch, + _FakeResponse(status=200, chunks=[b"image bytes"]), + ) + + result = await io.download_image_by_url( + "https://example.test/generate", + post=post, + post_data={"json": False} if post else None, + path=str(target_path), + ) + + assert result == str(target_path) + assert target_path.read_bytes() == b"image bytes"