Update

tests/unit/services/openai_gateway/gateway.py

@@ -331,6 +331,140 @@ def test_embeddings_disabled_returns_503(local_db, monkeypatch):
     assert resp.status_code == 503
 
 
+class FakeImageClient_openai_gateway:
+    def __init__(self, *a, **k):
+        self.calls = []
+
+    def build_request(self, method, url, headers=None, json=None, content=None):
+        return FakeRequest(method, url, json)
+
+    async def send(self, request):
+        self.calls.append((request.url, request.json_body))
+        body = request.json_body or {}
+        return FakeResp_openai_gateway(
+            payload={
+                "created": 1,
+                "model": body.get("model"),
+                "data": [{"b64_json": "aGVsbG8="}],
+                "usage": {"cost": 0.05},
+            }
+        )
+
+    async def aclose(self):
+        pass
+
+
+def test_images_remaps_alias_to_configured_model(local_db, monkeypatch):
+    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
+    svc = GatewayService()
+    cfg = svc.effective_config()
+    cfg["gateway_force_model"] = False
+    cfg["gateway_image_enabled"] = True
+    cfg["gateway_image_model"] = "black-forest-labs/flux.2-pro"
+    rt = svc.runtime()
+    run_async(
+        rt.handle_images(
+            {
+                "model": "molodetz-img-small",
+                "prompt": "award emblem",
+                "size": "512x512",
+            },
+            cfg,
+            ("guest", "test"),
+            "test",
+        )
+    )
+    assert rt._client.calls[-1][1]["model"] == "black-forest-labs/flux.2-pro"
+
+
+def test_image_route_overrides_upstream(local_db, monkeypatch):
+    from devplacepy.services.openai_gateway import routing
+
+    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
+    routing.provider_store.set(
+        routing.ProviderIn(
+            name="gwimg",
+            base_url="https://routed.example/v1/chat/completions",
+            api_key="img-key",
+        )
+    )
+    routing.model_store.set(
+        routing.ModelRouteIn(
+            source_model="custom-img",
+            provider="gwimg",
+            target_model="vendor/flux-pro",
+            kind="image",
+            price_input_per_m=0.06,
+        )
+    )
+    try:
+        svc = GatewayService()
+        cfg = svc.effective_config()
+        rt = svc.runtime()
+        run_async(
+            rt.handle_images(
+                {
+                    "model": "custom-img",
+                    "prompt": "trophy",
+                    "response_format": "b64_json",
+                },
+                cfg,
+                ("guest", "test"),
+                "test",
+            )
+        )
+        url, body = rt._client.calls[-1]
+        assert str(url) == "https://routed.example/v1/images"
+        assert body["model"] == "vendor/flux-pro"
+    finally:
+        routing.model_store.remove("custom-img")
+        routing.provider_store.remove("gwimg")
+
+
+def test_images_success_records_ledger(local_db, monkeypatch):
+    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
+    svc = GatewayService()
+    cfg = svc.effective_config()
+    cfg["gateway_image_enabled"] = True
+    cfg["gateway_image_model"] = "black-forest-labs/flux.2-pro"
+    rt = svc.runtime()
+    before = rt.image_calls
+    resp = run_async(
+        rt.handle_images(
+            {"model": "molodetz-img-small", "prompt": "badge"},
+            cfg,
+            ("guest", "img_ledger"),
+            "test",
+        )
+    )
+    assert resp.status_code == 200
+    payload = json.loads(bytes(resp.body).decode())
+    assert payload["data"][0]["b64_json"] == "aGVsbG8="
+    assert rt.image_calls == before + 1
+    row = get_table("gateway_usage_ledger").find_one(owner_id="img_ledger")
+    assert row is not None
+    assert row["backend"] == "image"
+    assert row["endpoint"] == "images/generations"
+    assert row["requested_model"] == "molodetz-img-small"
+    assert row["model"] == "black-forest-labs/flux.2-pro"
+    assert row["success"] == 1
+    assert float(row["cost_usd"]) == 0.05
+
+
+def test_images_disabled_returns_503(local_db, monkeypatch):
+    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
+    svc = GatewayService()
+    cfg = svc.effective_config()
+    cfg["gateway_image_enabled"] = False
+    rt = svc.runtime()
+    resp = run_async(
+        rt.handle_images(
+            {"prompt": "badge"}, cfg, ("guest", "test"), "test"
+        )
+    )
+    assert resp.status_code == 503
+
+
 def test_compute_cost_embed_branch():
     from devplacepy.services.openai_gateway.usage import (
         Pricing,