host
7a920cee59
首轮生产验收失败根因：_prepare_tenant 只把 IMAGE_GEN_* 透传 进 MCP env，没写 IMAGE_GEN_BASE_URL，也没带 SUB2API_BASE_URL， 于是 imagegen.py 回落到公网 sub2.diante.me，而它不在 Tinyproxy 白名单里，API 调用 403 Filtered。 修法：worker 在写 MCP env 前用 imagegen.build_mcp_env 算好 base（显式 IMAGE_GEN_BASE_URL > SUB2API_BASE_URL），并把 网关 host 加进 NO_PROXY；组装逻辑抽成可单测 helper，新增 2 项回归测试（本地网关解析、显式 base 优先）。
159 lines
6.7 KiB
Python
159 lines
6.7 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
|
|
from cloud import imagegen
|
|
from cloud.imagegen import ImageGenError, generate
|
|
from cloud.imagegen_mcp import handle_request
|
|
|
|
PNG_1x1 = (b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01"
|
|
b"\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde")
|
|
|
|
|
|
def _ok_request(body=None, key="k"):
|
|
def req(method, url, payload, token):
|
|
assert token == key
|
|
assert url.endswith("/v1/images/generations")
|
|
assert payload["model"] == "gpt-image-2.5-flare"
|
|
return 200, {"data": [body or {"url": "https://cdn.example/x.png",
|
|
"revised_prompt": "a red apple"}],
|
|
"usage": {"total_tokens": 1067}}
|
|
return req
|
|
|
|
|
|
class ImageGenTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self._old = dict(os.environ)
|
|
os.environ["IMAGE_GEN_API_KEY"] = "k"
|
|
os.environ["IMAGE_GEN_BASE_URL"] = "https://relay.example"
|
|
os.environ.pop("IMAGE_GEN_MODEL", None)
|
|
self._tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self._tmp.cleanup)
|
|
|
|
def tearDown(self):
|
|
os.environ.clear()
|
|
os.environ.update(self._old)
|
|
|
|
def test_url_download_writes_png_with_dimensions(self):
|
|
out = generate("a red apple", output_filename="apple.png",
|
|
request=_ok_request(),
|
|
download=lambda url: PNG_1x1,
|
|
cwd=self._tmp.name)
|
|
self.assertEqual(out["path"],
|
|
os.path.join(self._tmp.name, "apple.png"))
|
|
self.assertEqual(out["format"], "png")
|
|
self.assertEqual((out["width"], out["height"]), (1, 1))
|
|
self.assertEqual(out["bytes"], len(PNG_1x1))
|
|
with open(out["path"], "rb") as fh:
|
|
self.assertEqual(fh.read(), PNG_1x1)
|
|
|
|
def test_b64_response_decoded(self):
|
|
import base64
|
|
body = {"b64_json": base64.b64encode(PNG_1x1).decode()}
|
|
out = generate("a red apple", request=_ok_request(body),
|
|
cwd=self._tmp.name)
|
|
self.assertTrue(out["filename"].startswith("generated-"))
|
|
self.assertEqual(out["format"], "png")
|
|
|
|
def test_missing_key_is_configuration_error(self):
|
|
os.environ.pop("IMAGE_GEN_API_KEY")
|
|
with self.assertRaises(ImageGenError) as ctx:
|
|
generate("x", cwd=self._tmp.name)
|
|
self.assertIn("not configured", str(ctx.exception))
|
|
|
|
def test_path_traversal_rejected(self):
|
|
for bad in ("../evil.png", "sub/apple.png", "/abs/apple.png"):
|
|
with self.assertRaises(ImageGenError):
|
|
generate("x", output_filename=bad,
|
|
request=_ok_request(), cwd=self._tmp.name)
|
|
|
|
def test_bad_size_rejected_before_network(self):
|
|
def boom(*args, **kwargs):
|
|
raise AssertionError("must not call network")
|
|
with self.assertRaises(ImageGenError):
|
|
generate("x", size="999x999", request=boom,
|
|
cwd=self._tmp.name)
|
|
|
|
def test_relay_error_surfaced_without_key(self):
|
|
def req(method, url, payload, token):
|
|
return 404, {"error": {"message": "Model not supported",
|
|
"type": "model_not_found"}}
|
|
with self.assertRaises(ImageGenError) as ctx:
|
|
generate("x", request=req, cwd=self._tmp.name)
|
|
self.assertIn("404", str(ctx.exception))
|
|
self.assertNotIn("k", str(ctx.exception))
|
|
|
|
def test_mcp_tools_list_and_call(self):
|
|
listed = handle_request({"jsonrpc": "2.0", "id": 1,
|
|
"method": "tools/list"})
|
|
names = [t["name"] for t in listed["result"]["tools"]]
|
|
self.assertIn("image_generate", names)
|
|
called = handle_request(
|
|
{"jsonrpc": "2.0", "id": 2, "method": "tools/call",
|
|
"params": {"name": "image_generate",
|
|
"arguments": {"prompt": "a red apple"}}},
|
|
generate_fn=lambda *a, **k: {"path": "/w/apple.png"})
|
|
text = called["result"]["content"][0]["text"]
|
|
self.assertEqual(json.loads(text)["path"], "/w/apple.png")
|
|
self.assertFalse(called["result"]["isError"])
|
|
|
|
def test_mcp_unknown_tool(self):
|
|
resp = handle_request({"jsonrpc": "2.0", "id": 3,
|
|
"method": "tools/call",
|
|
"params": {"name": "nope",
|
|
"arguments": {}}})
|
|
self.assertEqual(resp["error"]["code"], -32602)
|
|
|
|
def test_build_mcp_env_resolves_local_gateway_base(self):
|
|
# Regression: first production verification failed 403
|
|
# Filtered because MCP env had no base URL and imagegen.py
|
|
# fell back to the public relay domain, which is not on
|
|
# the Tinyproxy whitelist. The local gateway origin must
|
|
# be resolved here and land in NO_PROXY.
|
|
env = imagegen.build_mcp_env(
|
|
{"SUB2API_BASE_URL": "http://172.17.0.1:28084",
|
|
"IMAGE_GEN_API_KEY": "k"},
|
|
{"HTTP_PROXY": "http://10.20.20.1:8888",
|
|
"NO_PROXY": "localhost,127.0.0.1"})
|
|
self.assertEqual(env["IMAGE_GEN_BASE_URL"],
|
|
"http://172.17.0.1:28084")
|
|
self.assertIn("172.17.0.1", env["NO_PROXY"])
|
|
self.assertEqual(env["IMAGE_GEN_API_KEY"], "k")
|
|
self.assertNotIn("SUB2API_API_KEY", env)
|
|
|
|
def test_build_mcp_env_explicit_base_wins(self):
|
|
env = imagegen.build_mcp_env(
|
|
{"IMAGE_GEN_BASE_URL": "https://relay.example",
|
|
"SUB2API_BASE_URL": "http://172.17.0.1:28084"})
|
|
self.assertEqual(env["IMAGE_GEN_BASE_URL"],
|
|
"https://relay.example")
|
|
self.assertIn("relay.example", env["NO_PROXY"])
|
|
|
|
def test_ensure_home_seeds_and_env(self):
|
|
home = os.path.join(self._tmp.name, "home")
|
|
os.makedirs(home)
|
|
with open(os.path.join(home, "config.toml"), "w") as fh:
|
|
fh.write("")
|
|
res = imagegen.ensure_codex_home_imagegen(
|
|
home, command_args=["python3", "/c/imagegen_mcp.py"])
|
|
self.assertTrue(res["config_changed"])
|
|
self.assertTrue(imagegen.ensure_imagegen_mcp_env(
|
|
home, {"IMAGE_GEN_API_KEY": "k",
|
|
"IMAGE_GEN_MODEL": "gpt-image-2.5-flare"}))
|
|
with open(os.path.join(home, "config.toml")) as fh:
|
|
text = fh.read()
|
|
self.assertIn("[mcp_servers.imagegen]", text)
|
|
self.assertIn("IMAGE_GEN_API_KEY", text)
|
|
# Idempotent: second env write with same values is a no-op.
|
|
self.assertFalse(imagegen.ensure_imagegen_mcp_env(
|
|
home, {"IMAGE_GEN_API_KEY": "k",
|
|
"IMAGE_GEN_MODEL": "gpt-image-2.5-flare"}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|