test_api.py

"""
Drive the generated REST endpoints against a running app.

This layer exists because `test_pipelines.py` calls `WorkflowExecutor` directly
and therefore skips gradio's *output component* postprocessing — which is
exactly where an image port's value shape matters. A bare `data:` URI passes
the executor happily and then dies in the endpoint with

    OSError: [Errno 22] Invalid argument: '...\\data:image\\png;base64,...'

because gradio treats the URI as a filename. Only these tests catch that.

    python apps/05_workflow1111/app.py                 # in another shell
    python apps/05_workflow1111/test_api.py [url]      # default 127.0.0.1:7865

Pass `--local` to exercise only the endpoints that need no Hugging Face token.
"""

import os
import sys
import time
import warnings

warnings.filterwarnings("ignore")
import logging  # noqa: E402

logging.disable(logging.CRITICAL)

for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(encoding="utf-8", errors="replace")
    except (AttributeError, ValueError):
        pass

from gradio_client import Client, handle_file  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
SAMPLES = os.path.join(HERE, "samples")

args = [a for a in sys.argv[1:] if not a.startswith("--")]
LOCAL_ONLY = "--local" in sys.argv
URL = args[0] if args else "http://127.0.0.1:7865"


def sample(name):
    path = os.path.join(SAMPLES, name)
    if not os.path.exists(path):
        raise SystemExit(f"missing sample {name} — run make_samples.py first")
    return handle_file(path)


# (api_name, args, needs_token)
CASES = [
    ("/annotator_map", lambda: [sample("control.jpg")], False),
    ("/png_info", lambda: [sample("with_parameters.png")], False),
    ("/upscaled_local", lambda: [sample("extras.jpg")], True),   # also hits 2 Spaces
    ("/image", lambda: ["a red fox in a snowy pine forest", "", "Cinematic",
                        "enhance fine detail"], True),
    ("/edited_image", lambda: [sample("init_image.jpg"),
                               "make it a snowy winter scene"], True),
    ("/generated_prompt", lambda: ["a lighthouse in a storm"], True),
    ("/recovered_prompt", lambda: [sample("interrogate.jpg")], True),
    ("/detected_objects", lambda: [sample("detect.jpg")], True),
    ("/x_y_grid", lambda: ["a lone tree on a hill",
                           "at sunrise | in a storm | at night | in fog"], True),
]


def describe(value):
    if isinstance(value, (list, tuple)):
        return " | ".join(describe(v) for v in value)
    if isinstance(value, dict):
        value = value.get("path") or value.get("url") or str(value)
    if isinstance(value, str) and os.path.isfile(value):
        try:
            from PIL import Image
            with Image.open(value) as im:
                return f"image {im.width}×{im.height} {im.format}"
        except Exception:
            return f"file ({os.path.getsize(value) // 1024} KB)"
    text = str(value).replace("\n", " ")
    return f"text: {text[:64]}"


def main():
    client = Client(URL, verbose=False)
    available = set(client.view_api(return_format="dict",
                                    print_info=False)["named_endpoints"])
    print(f"\n{URL} — {len(available)} endpoints\n")

    passed = failed = skipped = 0
    for name, build, needs_token in CASES:
        if name not in available:
            print(f"  skip  {name:20} (not exposed)")
            skipped += 1
            continue
        if LOCAL_ONLY and needs_token:
            skipped += 1
            continue
        t0 = time.time()
        try:
            result = client.predict(*build(), api_name=name)
        except Exception as e:
            failed += 1
            print(f"  FAIL  {name:20} ({time.time() - t0:6.1f}s)  "
                  f"{type(e).__name__}: {str(e)[:120]}")
        else:
            passed += 1
            print(f"  ok    {name:20} ({time.time() - t0:6.1f}s)  {describe(result)}")
        sys.stdout.flush()

    print(f"\n{'=' * 66}\n  {passed} passed, {failed} failed, {skipped} skipped\n{'=' * 66}")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())