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()