58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import unittest
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
|
import veida_client as kc # noqa: E402
|
|
|
|
|
|
class ClientTests(unittest.TestCase):
|
|
def test_wall_is_a_200(self):
|
|
with self.assertRaises(kc.VeidaError):
|
|
kc.parse_submit({"code": 0, "data": {"wall": True, "reason": "ip_daily"}})
|
|
|
|
def test_free_tier_polls_anon_query(self):
|
|
calls = []
|
|
|
|
def t(method, url, body, headers):
|
|
calls.append((method, url, body, headers))
|
|
if url.endswith("/api/ai/generate"):
|
|
return {"code": 0, "data": {"id": "t1"}}
|
|
return {"code": 0, "data": {"status": "success", "images": ["https://cdn/x.webp"]}}
|
|
|
|
url = kc.run("x", poll_every=0, transport=t, base="https://t", api_key="")
|
|
self.assertEqual(url, "https://cdn/x.webp")
|
|
self.assertIn("x-anon-id", calls[0][3])
|
|
self.assertEqual(json.loads(calls[0][2])["model"], "veida-image-v1")
|
|
self.assertIn("/api/ai/anon-query", calls[1][1])
|
|
|
|
def test_key_uses_account_and_prefers_clean(self):
|
|
calls = []
|
|
|
|
def t(method, url, body, headers):
|
|
calls.append((method, url, body, headers))
|
|
if url.endswith("/api/ai/generate"):
|
|
return {"code": 0, "data": {"id": "row1"}}
|
|
return {"code": 0, "data": {"status": "success", "images": ["https://cdn/m"], "cleanImages": ["https://cdn/c"]}}
|
|
|
|
url = kc.run("x", image_url="https://x/p.png", api_key="sk-abc", model="nano-banana-2",
|
|
poll_every=0, transport=t, base="https://t")
|
|
self.assertEqual(url, "https://cdn/c")
|
|
self.assertEqual(calls[0][3]["Authorization"], "Bearer sk-abc")
|
|
body = json.loads(calls[0][2])
|
|
self.assertEqual(body["model"], "nano-banana-2")
|
|
self.assertEqual(body["options"], {"image_input": ["https://x/p.png"]})
|
|
self.assertEqual(calls[1][1], "https://t/api/ai/query")
|
|
|
|
def test_upload_returns_first_url(self):
|
|
def t(method, url, body, headers):
|
|
self.assertIn(b"comfyui.png", body)
|
|
return {"code": 0, "data": {"urls": ["https://cdn/u.png"]}}
|
|
|
|
self.assertEqual(kc.upload_png(b"\x89PNG", transport=t, base="https://t"), "https://cdn/u.png")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|