Files
silveroxides-ComfyUI_Gemini…/tests/test_response_schema.py
T
silveroxides 8dbc870407 feat: add Gemini structured response support
Add a DOM schema builder and optional response formatting before image Autogrow. Return thoughts separately from final text and retain all outputs in schema-aware caches.
2026-09-11 19:06:22 +02:00

216 lines
8.9 KiB
Python

import asyncio
import json
import subprocess
from pathlib import Path
from types import SimpleNamespace
import pytest
from aiohttp import web
from middleware.cache_middleware import cache_control
from custom_nodes.ComfyUI_Gemini_Expanded_API import GeminiExtension
from custom_nodes.ComfyUI_Gemini_Expanded_API.response_schema import (
SSL_GeminiResponseSchema, compile_schema, default_entry, default_state, parse_state,
)
def add(state, parent, kind, name="", item=False):
identifier = f"n{state['next_id']}"
state["next_id"] += 1
entry = default_entry(kind)
entry["name"] = name
state["entries"][identifier] = entry
if item:
state["entries"][parent]["item"] = identifier
else:
state["entries"][parent]["fields"].append(identifier)
return identifier
def compile_state(state):
output = SSL_GeminiResponseSchema.execute(json.dumps(state))
assert output[0] == json.loads(output[1])
assert output.ui.as_dict() == {"text": (output[1],)}
return output[0]
def test_registration_and_dom_state_input():
assert SSL_GeminiResponseSchema in asyncio.run(GeminiExtension().get_node_list())
definition = SSL_GeminiResponseSchema.INPUT_TYPES()
assert list(definition["required"]) == ["schema_state"]
kind, metadata = definition["required"]["schema_state"]
assert kind == "GEMINI_SCHEMA_BUILDER"
assert metadata["socketless"] is True
assert metadata["state_version"] == 1
assert "templates" not in metadata
assert len(json.dumps(definition)) < 2000
assert compile_state(default_state()) == {"type": "object", "properties": {}}
def test_named_fields_order_and_required():
state = default_state()
add(state, "n0", "string", " title ")
second = add(state, "n0", "number", "score")
state["entries"][second]["required"] = False
assert compile_state(state) == {
"type": "object", "properties": {"title": {"type": "string"}, "score": {"type": "number"}}, "required": ["title"],
}
@pytest.mark.parametrize("count", [1, 9, 20])
def test_field_counts(count):
state = default_state()
for index in range(count):
add(state, "n0", "boolean", f"field_{index}")
result = compile_state(state)
assert list(result["properties"]) == [f"field_{index}" for index in range(count)]
assert all(schema == {"type": "boolean"} for schema in result["properties"].values())
@pytest.mark.parametrize("kinds", [["object", "object"], ["object", "array", "object"], ["array", "array"], ["object", "array"] * 4])
def test_nested_containers(kinds):
state = default_state()
state["entries"]["n0"]["type"] = kinds[0]
current = "n0"
expected = {"type": "string"}
for kind in reversed(kinds):
expected = {"type": "array", "items": expected} if kind == "array" else {
"type": "object", "properties": {"child": expected}, "required": ["child"],
}
for kind in kinds[1:] + ["string"]:
current = add(state, current, kind, "child", item=state["entries"][current]["type"] == "array")
assert compile_state(state) == expected
def test_disabled_and_inactive_state():
state = default_state()
field = add(state, "n0", "disabled")
child = add(state, field, "number")
state["entries"][child].update(minimum_mode="set", minimum="invalid")
assert compile_state(state) == {"type": "object", "properties": {}}
state["entries"][field].update(type="string", name="visible", choices="duplicate\nduplicate")
assert compile_state(state)["properties"] == {"visible": {"type": "string"}}
def test_unicode_names_and_description():
state = default_state()
state["entries"]["n0"]["description"] = " "
field = add(state, "n0", "string", "猫.[]")
state["entries"][field]["description"] = " Résumé\n続き "
result = compile_state(state)
assert "description" not in result
assert result["properties"]["猫.[]"]["description"] == " Résumé\n続き "
@pytest.mark.parametrize("kind", ["string", "integer", "number", "boolean"])
def test_root_types_and_constraints(kind):
state = default_state()
root = state["entries"]["n0"]
root["type"] = kind
expected = {"type": kind}
if kind == "string":
root.update(allowed_values="listed", choices=" positive\n\nnegative ")
expected["enum"] = ["positive", "negative"]
elif kind in ("integer", "number"):
root.update(minimum_mode="set", minimum=-2, maximum_mode="set", maximum=0)
expected.update(minimum=-2, maximum=0)
assert compile_state(state) == expected
@pytest.mark.parametrize("changes, path", [
({"type": "string", "allowed_values": "listed", "choices": ""}, "choices"),
({"type": "string", "allowed_values": "listed", "choices": "a\na"}, "choices"),
({"type": "number", "minimum_mode": "set", "minimum": True}, "minimum"),
({"type": "number", "minimum_mode": "set", "minimum": "1"}, "minimum"),
({"type": "number", "minimum_mode": "set", "minimum": float("inf")}, "minimum"),
({"type": "number", "maximum_mode": "set", "maximum": float("nan")}, "maximum"),
({"type": "integer", "minimum_mode": "set", "minimum": 0.5}, "minimum"),
({"type": "number", "minimum_mode": "set", "minimum": 2, "maximum_mode": "set", "maximum": 1}, "minimum"),
({"type": "array"}, "items"),
])
def test_invalid_active_settings(changes, path):
state = default_state()
state["entries"]["n0"].update(changes)
with pytest.raises(ValueError, match=path):
compile_state(state)
def test_name_errors_and_disabled_item():
state = default_state()
first = add(state, "n0", "string")
with pytest.raises(ValueError, match="name"):
compile_state(state)
state["entries"][first]["name"] = "same"
add(state, "n0", "number", "same")
with pytest.raises(ValueError, match="duplicate"):
compile_state(state)
state = default_state()
state["entries"]["n0"]["type"] = "array"
add(state, "n0", "disabled", item=True)
with pytest.raises(ValueError, match="disabled"):
compile_state(state)
@pytest.mark.parametrize("damage", ["root", "version", "cycle", "shared", "missing", "allocator", "orphan", "type", "mode"])
def test_invalid_structure(damage):
state = default_state()
field = add(state, "n0", "string", "name")
if damage == "root": state["root"] = "n99"
elif damage == "version": state["version"] = 2
elif damage == "cycle": state["entries"][field]["fields"] = ["n0"]
elif damage == "shared": state["entries"]["n0"]["fields"].append(field)
elif damage == "missing": state["entries"]["n0"]["fields"].append("n99")
elif damage == "allocator": state["next_id"] = 1
elif damage == "orphan": state["entries"]["n0"]["fields"] = []
elif damage == "type": state["entries"][field]["type"] = "invalid"
elif damage == "mode": state["entries"][field]["allowed_values"] = "invalid"
with pytest.raises(ValueError): parse_state(json.dumps(state))
def test_old_prototype_state_is_not_silently_replaced():
for value in ["object", "{", '{"type":"object"}', "[]"]:
with pytest.raises(ValueError): compile_schema(value)
@pytest.mark.parametrize("key, value", [("type", []), ("minimum_mode", []), ("item", 3), ("required", 1)])
def test_malformed_setting_shapes_report_validation_errors(key, value):
state = default_state()
state["entries"]["n0"][key] = value
with pytest.raises(ValueError):
compile_state(state)
@pytest.mark.parametrize("bound, value", [("minimum", 0), ("maximum", -1)])
def test_independent_numeric_bounds(bound, value):
state = default_state()
root = state["entries"]["n0"]
root.update(type="number", **{f"{bound}_mode": "set", bound: value})
assert compile_state(state) == {"type": "number", bound: value}
def test_browser_state_compiles_through_node():
module = (Path(__file__).resolve().parents[1] / "web" / "response_schema_model.js").as_uri()
script = f"""import {{ createState, addField, updateSetting, serializeState }} from {json.dumps(module)};
const state = createState();
const id = addField(state, state.root, 'string');
updateSetting(state, id, 'name', 'from_widgets');
console.log(serializeState(state));"""
result = subprocess.run(["node", "--input-type=module"], input=script, check=True, text=True, capture_output=True)
assert SSL_GeminiResponseSchema.execute(result.stdout)[0] == {
"type": "object", "properties": {"from_widgets": {"type": "string"}}, "required": ["from_widgets"],
}
def test_browser_modules_receive_comfy_no_store_headers():
directory = Path(__file__).resolve().parents[1] / "web"
assert not list(directory.glob("*.mjs"))
async def check():
async def handler(_request):
return web.Response(text="module")
for module in directory.glob("*.js"):
request = SimpleNamespace(path=f"/extensions/ComfyUI_Gemini_Expanded_API/{module.name}")
response = await cache_control(request, handler)
assert response.headers["Cache-Control"] == "no-store"
asyncio.run(check())