feat: use system curl for cmhub image downloads

This commit is contained in:
chengma
2026-07-08 09:26:25 +08:00
parent 4a499a17d1
commit cff61edce4
11 changed files with 355 additions and 9 deletions
+131
View File
@@ -93,6 +93,7 @@ class AITests(TempDirMixin, unittest.TestCase):
"title_alias": "title-standard",
"image_alias": "image-hd",
"connect_timeout": 3,
"download_with_curl": "false",
"check_balance_before_batch": False,
}
key_path = os.path.join(temp_dir, "cmhub.json")
@@ -637,6 +638,136 @@ class AITests(TempDirMixin, unittest.TestCase):
with self.assertRaises(ai.AIError):
ai._download_cmhub_image("https://cdn.example.com/a.png", 1, 1)
def test_cmhub_image_download_uses_curl_without_url_in_argv(self):
generated_png = self._png_bytes()
url = "https://cdn.example.com/generated.png?token=secret-token"
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
calls = []
def fake_run(args, **kwargs):
calls.append((args, kwargs))
self.assertIn("-K", args)
config_path = args[args.index("-K") + 1]
with open(config_path, "r", encoding="utf-8") as fh:
self.assertIn(url, fh.read())
self.assertNotIn(url, args)
self.assertIn("--connect-timeout", args)
self.assertEqual("3", args[args.index("--connect-timeout") + 1])
self.assertIn("--max-time", args)
self.assertEqual("650", args[args.index("--max-time") + 1])
self.assertIn("--max-filesize", args)
self.assertEqual(str(ai.CMHUB_IMAGE_MAX_BYTES), args[args.index("--max-filesize") + 1])
self.assertIn("--noproxy", args)
self.assertEqual("*", args[args.index("--noproxy") + 1])
self.assertFalse(kwargs["shell"])
output_path = args[args.index("--output") + 1]
with open(output_path, "wb") as fh:
fh.write(generated_png)
return SimpleNamespace(returncode=0, stdout=b"", stderr=b"")
with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
mock.patch("app.ai.subprocess.run", side_effect=fake_run), \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
image_bytes = ai._download_cmhub_image(
url,
connect_timeout=3,
read_timeout=650,
use_system_proxy=False,
download_with_curl="true",
)
self.assertEqual(generated_png, image_bytes)
self.assertEqual(1, len(calls))
def test_cmhub_image_download_skips_curl_for_private_url(self):
with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
mock.patch("app.ai.subprocess.run") as run:
with self.assertRaises(ai.AIError):
ai._download_cmhub_image(
"http://127.0.0.1/a.png",
connect_timeout=3,
read_timeout=650,
download_with_curl="true",
)
run.assert_not_called()
def test_cmhub_image_download_falls_back_to_requests_when_curl_fails(self):
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
with mock.patch("app.ai._find_system_curl", return_value=r"C:\Windows\System32\curl.exe"), \
mock.patch(
"app.ai.subprocess.run",
return_value=SimpleNamespace(returncode=28, stdout=b"", stderr=b"timeout"),
) as run, \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get) as get, \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
image_bytes = ai._download_cmhub_image(
"https://cdn.example.com/generated.png",
connect_timeout=3,
read_timeout=650,
download_with_curl="true",
)
self.assertEqual(generated_png, image_bytes)
self.assertEqual(1, run.call_count)
self.assertEqual(1, get.call_count)
def test_cmhub_image_download_auto_without_curl_uses_requests(self):
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
with mock.patch("app.ai.os.name", "nt"), \
mock.patch("app.ai._find_system_curl", return_value=""), \
mock.patch("app.ai.subprocess.run") as run, \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
image_bytes = ai._download_cmhub_image(
"https://cdn.example.com/generated.png",
connect_timeout=3,
read_timeout=650,
download_with_curl="auto",
)
self.assertEqual(generated_png, image_bytes)
run.assert_not_called()
def test_cmhub_image_download_auto_on_non_windows_uses_requests(self):
generated_png = self._png_bytes()
public_dns = [
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))
]
def fake_get(url, **kwargs):
return _RequestsResponse(content=generated_png)
with mock.patch("app.ai.os.name", "posix"), \
mock.patch("app.ai._find_system_curl", return_value="/usr/bin/curl"), \
mock.patch("app.ai.subprocess.run") as run, \
mock.patch.object(ai._cmhub_session(), "get", side_effect=fake_get), \
mock.patch("app.ai.socket.getaddrinfo", return_value=public_dns):
image_bytes = ai._download_cmhub_image(
"https://cdn.example.com/generated.png",
connect_timeout=3,
read_timeout=650,
download_with_curl="auto",
)
self.assertEqual(generated_png, image_bytes)
run.assert_not_called()
def test_cmhub_upstream_error_retries_and_keeps_metadata(self):
with self.make_temp_dir() as temp_dir:
cfg, key_path = self._cmhub_config(temp_dir)
+33
View File
@@ -101,6 +101,39 @@ class AppConfigTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_cmhub_download_with_curl_mode_is_normalized(self):
with self.make_temp_dir() as temp_dir:
config_path = os.path.join(temp_dir, "config.json")
config = appconfig.load_config(config_path)
self.assertEqual("auto", appconfig.cmhub_config(config)["download_with_curl"])
saved = appconfig.save_config(
{
"ai": {
"cmhub": {
"download_with_curl": True,
}
}
},
path=config_path,
)
self.assertEqual("true", appconfig.cmhub_config(saved)["download_with_curl"])
saved = appconfig.save_config(
{
"ai": {
"cmhub": {
"download_with_curl": "invalid",
}
}
},
path=config_path,
)
self.assertEqual("auto", appconfig.cmhub_config(saved)["download_with_curl"])
self.assert_removed(temp_dir)
def test_data_paths_resolve_under_default_data_dir(self):
with self.make_temp_dir() as temp_dir:
data_root = os.path.join(temp_dir, "data")