codex-app-server/cloud/tests/test_imagegen.py

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