236 lines
12 KiB
Python
236 lines
12 KiB
Python
"""真实浏览器 + 本地模拟上游的端到端验收，不访问付费服务。"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from unittest.mock import patch
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
from playwright.sync_api import sync_playwright, expect
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
PNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+a4l8AAAAASUVORK5CYII="
|
|
|
|
|
|
class Upstream(BaseHTTPRequestHandler):
|
|
def log_message(self, *args):
|
|
pass
|
|
|
|
def reply(self, body, status=200):
|
|
data = json.dumps(body).encode()
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(data)))
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
|
|
def do_GET(self):
|
|
self.reply({"data": [{"id": "vendor/test-image"}, {"id": "text-model"}]})
|
|
|
|
def do_POST(self):
|
|
payload = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
|
|
if payload["prompt"] == "fail-test":
|
|
self.reply({"error": "测试上游拒绝"}, 400)
|
|
else:
|
|
self.reply({"data": [{"b64_json": PNG} for _ in range(payload.get("n", 1))]})
|
|
|
|
|
|
class BrowserTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.runtime = tempfile.TemporaryDirectory(prefix="ai-image-browser-")
|
|
cls.upstream = ThreadingHTTPServer(("127.0.0.1", 0), Upstream)
|
|
threading.Thread(target=cls.upstream.serve_forever, daemon=True).start()
|
|
with socket.socket() as sock:
|
|
sock.bind(("127.0.0.1", 0))
|
|
port = sock.getsockname()[1]
|
|
cls.url = f"http://127.0.0.1:{port}"
|
|
cls.env = {**os.environ, "AI_IMAGE_DATA_DIR": str(Path(cls.runtime.name) / "data"),
|
|
"AI_IMAGE_DB": str(Path(cls.runtime.name) / "data" / "test.db"),
|
|
"AI_IMAGE_OUTPUT_DIR": str(Path(cls.runtime.name) / "output"),
|
|
"AI_IMAGE_SECRET_KEY": "browser-test-key", "AI_IMAGE_HOST": "127.0.0.1",
|
|
"NO_PROXY": "127.0.0.1,localhost"}
|
|
cls.server = subprocess.Popen([sys.executable, "main.py", "--port", str(port)], cwd=ROOT,
|
|
env=cls.env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
for _ in range(100):
|
|
try:
|
|
if httpx.get(cls.url + "/health", timeout=1, trust_env=False).is_success:
|
|
break
|
|
except httpx.HTTPError:
|
|
pass
|
|
time.sleep(.1)
|
|
else:
|
|
cls.server.terminate()
|
|
raise RuntimeError("测试服务启动失败")
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls.server.terminate()
|
|
cls.server.wait(timeout=10)
|
|
cls.upstream.shutdown()
|
|
cls.upstream.server_close()
|
|
cls.runtime.cleanup()
|
|
|
|
def test_browser_workflow(self):
|
|
with sync_playwright() as playwright:
|
|
browser = playwright.chromium.launch()
|
|
page = browser.new_page(viewport={"width": 1440, "height": 1000})
|
|
errors = []
|
|
page.on("pageerror", lambda error: errors.append(str(error)))
|
|
page.goto(self.url + "/providers.html")
|
|
page.locator("#fId").fill("browser")
|
|
page.locator("#fName").fill("浏览器测试供应商")
|
|
page.locator("#fBaseUrl").fill(f"http://127.0.0.1:{self.upstream.server_port}/v1")
|
|
page.locator("#fApiKey").fill("fake-key")
|
|
page.locator("#saveProviderBtn").click()
|
|
expect(page.locator("#providerList")).to_contain_text("浏览器测试供应商")
|
|
page.locator('[data-act="discover"]').click()
|
|
expect(page.locator("#discoverResult")).to_contain_text("vendor/test-image")
|
|
page.locator("#applyDiscoverBtn").click()
|
|
expect(page.locator("#toast-host")).to_contain_text("已保存 1 个模型")
|
|
|
|
page.goto(self.url + "/models.html")
|
|
expect(page.locator("#modelRows tr")).to_have_count(1)
|
|
page.locator("#disableAllBtn").click()
|
|
expect(page.locator("#statEnabled")).to_have_text("0")
|
|
page.locator("#enableAllBtn").click()
|
|
expect(page.locator("#statEnabled")).to_have_text("1")
|
|
|
|
page.goto(self.url + "/create.html")
|
|
expect(page.locator("#model")).to_have_value("vendor/test-image")
|
|
prompt = "a test image\n<script>window.promptInjected = true</script>"
|
|
page.locator("#prompt").fill(prompt)
|
|
page.locator("#count").select_option("2")
|
|
page.locator("#generateBtn").click()
|
|
expect(page.locator("#resultStatus")).to_contain_text("生成成功 · 2 张")
|
|
expect(page.locator("#resultGrid img")).to_have_count(2)
|
|
expect(page.locator("#recentGrid img")).to_have_count(2)
|
|
self.assertTrue(page.locator("#resultGrid img").first.evaluate("img => img.complete && img.naturalWidth > 0"))
|
|
page.locator("#prompt").fill("fail-test")
|
|
page.locator("#generateBtn").click()
|
|
expect(page.locator("#resultStatus")).to_have_text("生成失败")
|
|
expect(page.locator("#generateBtn")).not_to_have_class("is-loading")
|
|
|
|
page.goto(self.url + "/records.html")
|
|
expect(page.locator("#recRows tr")).to_have_count(2)
|
|
page.locator('[data-status="failed"]').click()
|
|
expect(page.locator("#recRows tr")).to_have_count(1)
|
|
page.locator("#recRows tr").click()
|
|
expect(page.locator("#detailModal")).to_have_class("modal-mask open")
|
|
page.keyboard.press("Escape")
|
|
|
|
page.goto(self.url + "/index.html")
|
|
expect(page.locator(".folder-card")).to_have_count(1)
|
|
page.locator(".folder-card").click()
|
|
expect(page.locator("#galleryGrid img")).to_have_count(2)
|
|
page.locator("#galleryGrid .gal-item").first.click()
|
|
expect(page.locator("#lbPrompt")).to_be_visible()
|
|
expect(page.locator("#lbPromptText")).to_have_text(prompt)
|
|
self.assertFalse(page.evaluate("window.promptInjected === true"))
|
|
page.on("dialog", lambda dialog: dialog.accept())
|
|
page.locator("#lbDelete").click()
|
|
expect(page.locator("#galleryGrid img")).to_have_count(1)
|
|
expect(page.locator("#statTotal")).to_have_text("1")
|
|
|
|
page.goto(self.url + "/skills.html")
|
|
expect(page.locator("#mdView")).to_contain_text("调用流程")
|
|
page.locator("#rawBtn").click()
|
|
expect(page.locator("#rawView")).to_be_visible()
|
|
for width in (2560, 1440, 390):
|
|
page.set_viewport_size({"width": width, "height": 900})
|
|
for name in ("index", "create", "providers", "models", "records", "skills"):
|
|
page.goto(self.url + f"/{name}.html")
|
|
expect(page.locator(".top-bar h1")).to_be_visible()
|
|
page.wait_for_load_state("networkidle")
|
|
overflow = page.evaluate("document.documentElement.scrollWidth > innerWidth")
|
|
self.assertFalse(overflow, f"{name} 页面在 {width}px 出现横向溢出")
|
|
content = page.locator(".app-content").bounding_box()
|
|
body = page.locator(".page-body").bounding_box()
|
|
self.assertAlmostEqual(body["width"], content["width"], delta=1,
|
|
msg=f"{name} 页面在 {width}px 未铺满内容区域")
|
|
|
|
page.goto(self.url + "/providers.html")
|
|
page.locator('[data-act="del"]').click()
|
|
page.locator("#confirmDeleteBtn").click()
|
|
expect(page.locator("#providerList")).not_to_contain_text("浏览器测试供应商")
|
|
self.assertEqual(errors, [])
|
|
legacy = Path(self.runtime.name) / "output" / "legacy-preview"
|
|
legacy.mkdir(exist_ok=True)
|
|
(legacy / "legacy.png").write_bytes(base64.b64decode(PNG))
|
|
page.goto(self.url + "/index.html")
|
|
page.locator(".folder-card").filter(has_text="legacy-preview").click()
|
|
page.locator("#galleryGrid .gal-item").first.click()
|
|
expect(page.locator("#lightbox")).to_have_class("lightbox open")
|
|
expect(page.locator("#lbPrompt")).to_be_hidden()
|
|
expect(page.locator("#lbPromptText")).to_be_empty()
|
|
browser.close()
|
|
|
|
def test_provider_state_and_default(self):
|
|
with sync_playwright() as playwright:
|
|
browser = playwright.chromium.launch()
|
|
page = browser.new_page()
|
|
page.goto(self.url + "/providers.html")
|
|
page.locator("#fId").fill("state")
|
|
page.locator("#fName").fill("状态测试")
|
|
page.locator("#fBaseUrl").fill(f"http://127.0.0.1:{self.upstream.server_port}/v1")
|
|
page.locator("#fApiKey").fill("fake-key")
|
|
page.locator("#fDefaultModel").fill("z-image")
|
|
page.locator("#saveProviderBtn").click()
|
|
expect(page.locator('[data-act="edit"]')).to_be_visible()
|
|
response = page.request.post(self.url + "/api/providers/state/models",
|
|
data={"model_ids": ["a-image", "z-image"]})
|
|
self.assertEqual(response.status, 200)
|
|
page.goto(self.url + "/create.html")
|
|
expect(page.locator("#model")).to_have_value("z-image")
|
|
page.request.patch(self.url + "/api/models/state%3Az-image", data={"enabled": False})
|
|
page.reload()
|
|
expect(page.locator("#model")).to_have_value("a-image")
|
|
page.goto(self.url + "/providers.html")
|
|
page.locator('[data-act="edit"]').click()
|
|
page.locator("#fEnabled").uncheck()
|
|
page.locator("#saveProviderBtn").click()
|
|
expect(page.locator("#providerList")).to_contain_text("停用")
|
|
page.goto(self.url + "/create.html")
|
|
expect(page.locator("#provider")).to_have_value("")
|
|
expect(page.locator("#model")).to_have_value("")
|
|
page.goto(self.url + "/providers.html")
|
|
page.locator('[data-act="edit"]').click()
|
|
expect(page.locator("#fEnabled")).not_to_be_checked()
|
|
page.locator("#fEnabled").check()
|
|
page.locator("#saveProviderBtn").click()
|
|
expect(page.locator("#providerList")).to_contain_text("启用")
|
|
page.request.delete(self.url + "/api/providers/state")
|
|
browser.close()
|
|
|
|
def test_mcp_http_and_stdio(self):
|
|
from fastmcp import Client
|
|
from fastmcp.client.transports import StdioTransport
|
|
|
|
async def check(transport):
|
|
async with Client(transport) as client:
|
|
names = {tool.name for tool in await client.list_tools()}
|
|
self.assertEqual(names, {"list_providers", "discover_models", "generate_image",
|
|
"list_generations", "get_skill_doc"})
|
|
result = await client.call_tool("get_skill_doc", {})
|
|
self.assertFalse(result.is_error)
|
|
self.assertIn("ai-image", str(result))
|
|
|
|
with patch.dict(os.environ, {"NO_PROXY": "127.0.0.1,localhost"}):
|
|
asyncio.run(check(self.url + "/mcp"))
|
|
asyncio.run(check(StdioTransport(command=sys.executable, args=[str(ROOT / "main.py"), "--stdio"],
|
|
env=self.env, cwd=str(ROOT))))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|