Files
ai-image/tests/test_browser.py
T

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()