849 lines
31 KiB
Python
849 lines
31 KiB
Python
"""覆盖率 85% 第 4 步第 5 批：workers / autoscaler / playwright 长尾。
|
||
|
||
只动测试。workers 与 autoscaler 纯函数走真实 SQLite；
|
||
playwright Browser 走类自带的 browser_factory 注入点（假
|
||
playwright 对象），ensure_* 只写临时 codex_home，绝不启动
|
||
真实浏览器。autoscaler 的进程操作以 Popen 替身覆盖启动，
|
||
不拉真实 worker 子进程。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from unittest import mock
|
||
|
||
from cloud import autoscaler as autoscaler_mod
|
||
from cloud import playwright_browser as pw_mod
|
||
from cloud import workers as workers_mod
|
||
|
||
|
||
def _sqlite_db(path, tasks=True, workers=False):
|
||
conn = sqlite3.connect(path)
|
||
if tasks:
|
||
conn.execute(
|
||
"CREATE TABLE tasks(id TEXT PRIMARY KEY, status TEXT,"
|
||
" depends_on TEXT, executor TEXT, coordinator_parts TEXT,"
|
||
" last_heartbeat REAL)")
|
||
if workers:
|
||
conn.execute(
|
||
"CREATE TABLE workers(worker_id TEXT PRIMARY KEY,"
|
||
" executor TEXT, capacity INTEGER, running INTEGER,"
|
||
" last_seen REAL, created_at REAL)")
|
||
conn.commit()
|
||
return conn
|
||
|
||
|
||
class WorkersTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db = os.path.join(self.tmp.name, "w.db")
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def test_threshold_override_and_env(self):
|
||
self.assertEqual(workers_mod.alive_threshold_seconds(30), 30)
|
||
self.assertEqual(workers_mod.alive_threshold_seconds("bogus"),
|
||
120)
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_WORKER_ALIVE_SECONDS": "bogus"}):
|
||
self.assertEqual(workers_mod.alive_threshold_seconds(), 120)
|
||
with mock.patch.dict(os.environ,
|
||
{"CLOUD_WORKER_ALIVE_SECONDS": "45"}):
|
||
self.assertEqual(workers_mod.alive_threshold_seconds(), 45)
|
||
|
||
def test_register_and_heartbeat_flows(self):
|
||
with self.assertRaises(ValueError):
|
||
workers_mod.register_worker(self.db, "")
|
||
with self.assertRaises(ValueError):
|
||
workers_mod.heartbeat_worker(self.db, None)
|
||
row = workers_mod.register_worker(self.db, "w1",
|
||
executor="docker", capacity=2,
|
||
running=1)
|
||
self.assertEqual(row["worker_id"], "w1")
|
||
self.assertEqual(row["executor"], "docker")
|
||
self.assertEqual(row["capacity"], 2)
|
||
self.assertEqual(row["running"], 1)
|
||
row = workers_mod.register_worker(self.db, "w1",
|
||
executor="local", capacity=3)
|
||
self.assertEqual(row["executor"], "local")
|
||
self.assertEqual(row["running"], 0)
|
||
# 心跳未给 executor/capacity 时保留旧值，只刷时间。
|
||
row = workers_mod.heartbeat_worker(self.db, "w1", running=2)
|
||
self.assertEqual(row["executor"], "local")
|
||
self.assertEqual(row["capacity"], 3)
|
||
self.assertEqual(row["running"], 2)
|
||
# 未注册 worker 的心跳即注册。
|
||
row = workers_mod.heartbeat_worker(self.db, "w2")
|
||
self.assertEqual(row["capacity"], 1)
|
||
self.assertEqual(row["executor"], None)
|
||
|
||
def test_list_alive_and_stale(self):
|
||
workers_mod.register_worker(self.db, "fresh")
|
||
workers_mod.register_worker(self.db, "stale")
|
||
conn = sqlite3.connect(self.db)
|
||
conn.execute("UPDATE workers SET last_seen = ?"
|
||
" WHERE worker_id = 'stale'",
|
||
(time.time() - 10000,))
|
||
conn.commit()
|
||
conn.close()
|
||
rows = workers_mod.list_workers(self.db, alive_seconds=120)
|
||
by_id = {r["worker_id"]: r for r in rows}
|
||
self.assertTrue(by_id["fresh"]["alive"])
|
||
self.assertFalse(by_id["stale"]["alive"])
|
||
self.assertEqual([r["worker_id"] for r in rows],
|
||
["fresh", "stale"])
|
||
|
||
def test_ensure_old_table_adds_running_column(self):
|
||
conn = sqlite3.connect(self.db)
|
||
conn.execute(
|
||
"CREATE TABLE workers(worker_id TEXT PRIMARY KEY,"
|
||
" executor TEXT, capacity INTEGER, last_seen REAL,"
|
||
" created_at REAL)")
|
||
conn.execute("INSERT INTO workers VALUES('old', 'docker',"
|
||
" 1, 1000, 1000)")
|
||
conn.commit()
|
||
conn.close()
|
||
workers_mod.ensure_workers_table(self.db)
|
||
conn = sqlite3.connect(self.db)
|
||
cols = [r[1] for r in
|
||
conn.execute("PRAGMA table_info(workers)").fetchall()]
|
||
conn.close()
|
||
self.assertIn("running", cols)
|
||
# 幂等：已含 running 列时再跑不报错。
|
||
workers_mod.ensure_workers_table(self.db)
|
||
|
||
def test_worker_to_dict_missing_keys(self):
|
||
d = workers_mod._worker_to_dict({"worker_id": "w"})
|
||
self.assertEqual(d["worker_id"], "w")
|
||
self.assertIsNone(d["executor"])
|
||
self.assertEqual(d["running"], 0)
|
||
d = workers_mod._worker_to_dict(None)
|
||
self.assertIsNone(d["worker_id"])
|
||
self.assertEqual(d["running"], 0)
|
||
|
||
|
||
class _FakeCursor:
|
||
def __init__(self, rows=None):
|
||
self.rows = rows or []
|
||
|
||
def fetchone(self):
|
||
return self.rows[0] if self.rows else None
|
||
|
||
def fetchall(self):
|
||
return self.rows
|
||
|
||
|
||
class _FakeConn:
|
||
"""workers 依赖桩：PRAGMA 抛错、information_schema 可控。"""
|
||
|
||
def __init__(self, info_row=None, count_row=None):
|
||
self.info_row = info_row
|
||
self.count_row = count_row
|
||
self.executed = []
|
||
|
||
def execute(self, sql, params=None):
|
||
self.executed.append(sql)
|
||
if "PRAGMA table_info" in sql:
|
||
raise RuntimeError("no pragma")
|
||
if "information_schema.columns" in sql and "COUNT" in sql:
|
||
return _FakeCursor([self.count_row] if self.count_row
|
||
else [])
|
||
if "information_schema.columns" in sql:
|
||
if self.info_row is None:
|
||
raise RuntimeError("no information_schema")
|
||
return _FakeCursor([self.info_row])
|
||
if "ALTER TABLE" in sql:
|
||
raise RuntimeError("alter not supported here")
|
||
return _FakeCursor()
|
||
|
||
def commit(self):
|
||
pass
|
||
|
||
def rollback(self):
|
||
pass
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
|
||
class WorkersFakeConnTest(unittest.TestCase):
|
||
def test_ensure_running_column_pg_fallback(self):
|
||
conn = _FakeConn(count_row=(0,))
|
||
workers_mod._ensure_running_column(conn)
|
||
self.assertTrue(any("ADD COLUMN IF NOT EXISTS" in s
|
||
for s in conn.executed))
|
||
|
||
def test_ensure_running_column_pg_query_fails(self):
|
||
conn = _FakeConn()
|
||
# information_schema 两次查询都抛错：函数只回滚不抛。
|
||
workers_mod._ensure_running_column(conn)
|
||
|
||
def test_ensure_table_pg_real_column_path(self):
|
||
conn = _FakeConn(info_row=("real",))
|
||
with mock.patch.object(workers_mod, "_db_connect",
|
||
return_value=conn):
|
||
with mock.patch.object(workers_mod,
|
||
"_ensure_running_column"):
|
||
workers_mod.ensure_workers_table("fake.db")
|
||
self.assertTrue(any("ALTER TABLE workers ALTER COLUMN" in s
|
||
for s in conn.executed))
|
||
|
||
|
||
class AutoscalerPureTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db = os.path.join(self.tmp.name, "a.db")
|
||
self.conn = _sqlite_db(self.db, workers=True)
|
||
|
||
def tearDown(self):
|
||
self.conn.close()
|
||
self.tmp.cleanup()
|
||
|
||
def _add_task(self, tid, status, deps=None, executor=None,
|
||
parts=None, hb=None):
|
||
self.conn.execute(
|
||
"INSERT INTO tasks VALUES(?, ?, ?, ?, ?, ?)",
|
||
(tid, status,
|
||
json.dumps(deps) if deps is not None else None,
|
||
executor, parts, hb))
|
||
self.conn.commit()
|
||
|
||
def test_queue_depth_variants(self):
|
||
self._add_task("a", "queued")
|
||
self._add_task("b", "queued", deps=["a"])
|
||
self._add_task("c", "completed")
|
||
self._add_task("d", "queued", deps=["c"])
|
||
self._add_task("e", "queued", deps=None)
|
||
self.assertEqual(autoscaler_mod.queue_depth(self.db), 3)
|
||
# 坏 JSON 依赖当无依赖。
|
||
self.conn.execute(
|
||
"INSERT INTO tasks VALUES('f', 'queued', 'not-json',"
|
||
" NULL, NULL, NULL)")
|
||
self.conn.commit()
|
||
self.assertEqual(autoscaler_mod.queue_depth(self.db), 4)
|
||
self._add_task("i", "queued", executor="docker")
|
||
self.assertEqual(autoscaler_mod.queue_depth(
|
||
self.db, executor="local"), 4)
|
||
self.assertEqual(autoscaler_mod.queue_depth(
|
||
self.db, executor="docker"), 5)
|
||
# 成环任务永不就绪（依赖不在 completed 集里），但不影响
|
||
# 其他就绪任务的计数：仍为 a、d、e、f、i 共 5。
|
||
self.conn.execute(
|
||
"INSERT INTO tasks VALUES('g', 'queued', '[\"h\"]',"
|
||
" NULL, NULL, NULL)")
|
||
self.conn.execute(
|
||
"INSERT INTO tasks VALUES('h', 'queued', '[\"g\"]',"
|
||
" NULL, NULL, NULL)")
|
||
self.conn.commit()
|
||
self.assertEqual(autoscaler_mod.queue_depth(self.db), 5)
|
||
|
||
def test_desired_workers(self):
|
||
f = autoscaler_mod.desired_workers
|
||
self.assertEqual(f(0, 2, 1, 3), 1)
|
||
self.assertEqual(f(10, 0, 1, 3), 3)
|
||
self.assertEqual(f(10, 3, 1, 2), 2)
|
||
self.assertEqual(f(5, 2, 1, 10), 3)
|
||
self.assertEqual(f(-1, 2, 2, 5), 2)
|
||
|
||
def test_live_heartbeat(self):
|
||
f = autoscaler_mod._is_live_heartbeat
|
||
now = time.time()
|
||
self.assertFalse(f(None, now, 120))
|
||
self.assertFalse(f("bogus", now, 120))
|
||
self.assertTrue(f(now - 10, now, 120))
|
||
self.assertTrue(f(now + 60, now, 120))
|
||
self.assertFalse(f(now - 1000, now, 120))
|
||
|
||
def test_running_task_helpers(self):
|
||
now = time.time()
|
||
self._add_task("r1", "running",
|
||
parts='[{"name": "a", "prompt": "do A"}]',
|
||
hb=now - 5)
|
||
self._add_task("r2", "running", parts="bad-json",
|
||
hb=now - 5)
|
||
self._add_task("r3", "running",
|
||
parts='[{"name": "a", "prompt": "do A"}]',
|
||
hb=now - 10000)
|
||
self.assertTrue(autoscaler_mod._has_live_running_coordinator(
|
||
self.db, now=now))
|
||
self.assertTrue(autoscaler_mod._has_live_running_task(
|
||
self.db, now=now))
|
||
# 全部心跳陈旧时两者皆 False。
|
||
self.conn.execute("UPDATE tasks SET last_heartbeat = ?",
|
||
(now - 10000,))
|
||
self.conn.commit()
|
||
self.assertFalse(autoscaler_mod._has_live_running_coordinator(
|
||
self.db, now=now))
|
||
self.assertFalse(autoscaler_mod._has_live_running_task(
|
||
self.db, now=now))
|
||
|
||
def test_worker_registry_running(self):
|
||
now = time.time()
|
||
f = autoscaler_mod._worker_registry_running
|
||
self.assertEqual(f(self.db, "nope", now=now), 0)
|
||
self.conn.execute(
|
||
"INSERT INTO workers VALUES('w1', 'docker', 2, 3, ?, ?)",
|
||
(now - 5, now - 5))
|
||
self.conn.execute(
|
||
"INSERT INTO workers VALUES('w2', 'docker', 2, 9, ?, ?)",
|
||
(now - 10000, now - 10000))
|
||
self.conn.execute(
|
||
"INSERT INTO workers VALUES('w3', 'docker', 2, NULL, ?, ?)",
|
||
(now - 5, now - 5))
|
||
self.conn.commit()
|
||
self.assertEqual(f(self.db, "w1", now=now), 3)
|
||
self.assertEqual(f(self.db, "w2", now=now), 0)
|
||
self.assertEqual(f(self.db, "w3", now=now), 0)
|
||
# workers 表缺失时按 0 处理。
|
||
os.remove(self.db)
|
||
self.assertEqual(f(self.db, "w1", now=now), 0)
|
||
|
||
def test_victim_is_busy(self):
|
||
now = time.time()
|
||
self.conn.execute(
|
||
"INSERT INTO workers VALUES('w1', 'docker', 2, 1, ?, ?)",
|
||
(now - 5, now - 5))
|
||
self.conn.commit()
|
||
# 注册表报 busy 但 DB 无活任务：幽灵占位，不拦。
|
||
self.assertFalse(autoscaler_mod._victim_is_busy(
|
||
self.db, "w1", now=now))
|
||
self._add_task("r1", "running", hb=now - 5)
|
||
self.assertTrue(autoscaler_mod._victim_is_busy(
|
||
self.db, "w1", now=now))
|
||
# 注册表报 0：直接不忙。
|
||
self.assertFalse(autoscaler_mod._victim_is_busy(
|
||
self.db, "nobody", now=now))
|
||
|
||
|
||
class _FakeProc:
|
||
def __init__(self, pid=4242, poll_result=None, wait_raises=False):
|
||
self.pid = pid
|
||
self._poll = poll_result
|
||
self.returncode = poll_result
|
||
self.signals = []
|
||
self.wait_calls = 0
|
||
self.wait_raises = wait_raises
|
||
|
||
def poll(self):
|
||
return self._poll
|
||
|
||
def send_signal(self, sig):
|
||
self.signals.append(sig)
|
||
|
||
def wait(self, timeout=None):
|
||
self.wait_calls += 1
|
||
if self.wait_raises and self.wait_calls == 1:
|
||
import subprocess as _sp
|
||
raise _sp.TimeoutExpired("cmd", timeout)
|
||
self._poll = 0
|
||
self.returncode = 0
|
||
return 0
|
||
|
||
|
||
class AutoscalerProcTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db = os.path.join(self.tmp.name, "a.db")
|
||
self.state = os.path.join(self.tmp.name, "state.json")
|
||
self.log_dir = os.path.join(self.tmp.name, "logs")
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _scaler(self, **kw):
|
||
kw.setdefault("worker_args", [])
|
||
kw.setdefault("state_file", self.state)
|
||
kw.setdefault("log_dir", self.log_dir)
|
||
return autoscaler_mod.Autoscaler(self.db, **kw)
|
||
|
||
def test_worker_command_variants(self):
|
||
s = self._scaler(worker_cmd=["/bin/worker", "--x"])
|
||
self.assertEqual(s._worker_command("w1"),
|
||
["/bin/worker", "--x"])
|
||
s = self._scaler(worker_args=["--loop", "--db", "/tmp/d.db",
|
||
"--host-name", "h1"])
|
||
cmd = s._worker_command("w9")
|
||
self.assertIn("--worker-id", cmd)
|
||
self.assertEqual(cmd.count("--loop"), 1)
|
||
self.assertEqual(cmd.count("--db"), 1)
|
||
self.assertEqual(cmd.count("--host-name"), 1)
|
||
s = self._scaler(worker_args=[])
|
||
cmd = s._worker_command("w1")
|
||
self.assertIn("--loop", cmd)
|
||
self.assertIn("--db", cmd)
|
||
self.assertIn("--host-name", cmd)
|
||
|
||
def test_start_stop_reap_next_id(self):
|
||
s = self._scaler()
|
||
procs = []
|
||
with mock.patch.object(autoscaler_mod.subprocess, "Popen",
|
||
side_effect=lambda *a, **k: procs.pop(0)):
|
||
procs.append(_FakeProc(pid=111))
|
||
p = s._start_worker("autoscale-1")
|
||
self.assertEqual(p.pid, 111)
|
||
self.assertTrue(os.path.exists(
|
||
os.path.join(self.log_dir, "autoscale-1.log")))
|
||
self.assertEqual(s._next_worker_id(), "autoscale-2")
|
||
# 停不存在的 worker 返回 True。
|
||
self.assertTrue(s._stop_worker("ghost"))
|
||
# 超时路径：wait 先 TimeoutExpired，再 killpg + 等到退出。
|
||
s.workers["autoscale-1"]["proc"] = _FakeProc(
|
||
pid=111, wait_raises=True)
|
||
with mock.patch.object(autoscaler_mod.os, "killpg"):
|
||
self.assertTrue(s._stop_worker("autoscale-1"))
|
||
self.assertNotIn("autoscale-1", s.workers)
|
||
# reap：已死 proc 被回收、日志句柄关闭。
|
||
dead = _FakeProc(pid=222, poll_result=1)
|
||
fh = open(os.path.join(self.tmp.name, "x.log"), "wb")
|
||
s.workers["autoscale-2"] = {"proc": dead, "started_at": 1.0,
|
||
"log": fh}
|
||
s._reap()
|
||
self.assertNotIn("autoscale-2", s.workers)
|
||
self.assertTrue(fh.closed)
|
||
|
||
def test_write_state_and_run_loop_stop(self):
|
||
s = self._scaler()
|
||
s.workers["autoscale-1"] = {
|
||
"proc": _FakeProc(pid=777), "started_at": 123.0,
|
||
"log": None}
|
||
s.last_depth = 4
|
||
s.last_desired = 2
|
||
s._write_state()
|
||
with open(self.state) as f:
|
||
state = json.load(f)
|
||
self.assertEqual(state["actual"], 1)
|
||
self.assertEqual(state["workers"][0]["pid"], 777)
|
||
self.assertEqual(state["depth"], 4)
|
||
# 无 state_file 时 _write_state 是空操作。
|
||
s2 = autoscaler_mod.Autoscaler(self.db, [])
|
||
s2._write_state()
|
||
# run_loop：stop_event 已置位则立即退出；finally 停掉
|
||
# 全部 worker 后再写 state，actual 回到 0。
|
||
import threading
|
||
ev = threading.Event()
|
||
ev.set()
|
||
s.run_loop(stop_event=ev)
|
||
with open(self.state) as f:
|
||
state = json.load(f)
|
||
self.assertEqual(state["actual"], 0)
|
||
s.request_stop()
|
||
self.assertTrue(s._stop)
|
||
|
||
|
||
# ---------------------------------------------------------- playwright
|
||
|
||
class _SimpleQueue:
|
||
"""queue.Queue 的最小替身（只给 screencast sink 用）。
|
||
|
||
刻意不用 stdlib queue：cloud/tests/test_queue_sandbox.py
|
||
靠 cloud/queue.py 遮蔽 stdlib queue（它把 cloud/ 插进
|
||
sys.path 再 from queue import）；discover 按字母顺序先
|
||
导入本文件时若已缓存 stdlib queue，那个模块就 ImportError。
|
||
本步实测：import queue 一行导致 cloud 套件 47 项未跑。
|
||
"""
|
||
|
||
def __init__(self):
|
||
self._items = []
|
||
|
||
def put(self, item):
|
||
self._items.append(item)
|
||
|
||
def get_nowait(self):
|
||
if not self._items:
|
||
raise RuntimeError("empty")
|
||
return self._items.pop(0)
|
||
|
||
|
||
class _FakeMouse:
|
||
def __init__(self):
|
||
self.moves = []
|
||
self.clicks = []
|
||
|
||
def move(self, x, y):
|
||
self.moves.append((x, y))
|
||
|
||
def click(self, x, y):
|
||
self.clicks.append((x, y))
|
||
|
||
|
||
class _FakeKeyboard:
|
||
def __init__(self):
|
||
self.pressed = []
|
||
self.typed = []
|
||
|
||
def press(self, key):
|
||
self.pressed.append(key)
|
||
|
||
def type(self, text):
|
||
self.typed.append(text)
|
||
|
||
|
||
class _FakeLocator:
|
||
def __init__(self, page, selector):
|
||
self.page = page
|
||
self.selector = selector
|
||
self.clicked = False
|
||
|
||
def inner_text(self, timeout=None):
|
||
return self.page.body_text
|
||
|
||
@property
|
||
def first(self):
|
||
return self
|
||
|
||
def click(self, timeout=None):
|
||
self.clicked = True
|
||
self.page.locator_clicks.append(self.selector)
|
||
|
||
|
||
class _FakeCDP:
|
||
def __init__(self):
|
||
self.handlers = {}
|
||
self.sent = []
|
||
|
||
def on(self, event, handler):
|
||
self.handlers[event] = handler
|
||
|
||
def send(self, method, params=None):
|
||
self.sent.append((method, params))
|
||
if method == "Page.captureScreenshot":
|
||
return {"data": "eager-b64"}
|
||
return {}
|
||
|
||
|
||
class _FakeContext:
|
||
def __init__(self, page):
|
||
self.page = page
|
||
self.cdp = None
|
||
|
||
def new_cdp_session(self, page):
|
||
self.cdp = _FakeCDP()
|
||
return self.cdp
|
||
|
||
|
||
class _FakePage:
|
||
def __init__(self):
|
||
self.url = "about:blank"
|
||
self._title = "Fake Title"
|
||
self.body_text = "hello world"
|
||
self.mouse = _FakeMouse()
|
||
self.keyboard = _FakeKeyboard()
|
||
self.locator_clicks = []
|
||
self.context = _FakeContext(self)
|
||
self.shots = []
|
||
self.goto_fail = False
|
||
self.eval_result = None
|
||
|
||
def title(self):
|
||
return self._title
|
||
|
||
def goto(self, url, timeout=None, wait_until=None):
|
||
if self.goto_fail:
|
||
raise RuntimeError("net fail")
|
||
self.url = url
|
||
|
||
def click(self, selector, timeout=None):
|
||
self.locator_clicks.append(selector)
|
||
|
||
def fill(self, selector, text, timeout=None):
|
||
self.filled = (selector, text)
|
||
|
||
def locator(self, selector):
|
||
return _FakeLocator(self, selector)
|
||
|
||
def evaluate(self, script):
|
||
if self.eval_result is not None:
|
||
return self.eval_result
|
||
return "eval-ok"
|
||
|
||
def screenshot(self, path=None):
|
||
with open(path, "wb") as f:
|
||
f.write(b"PNGDATA")
|
||
self.shots.append(path)
|
||
|
||
def close(self):
|
||
self.closed = True
|
||
|
||
|
||
class _FakeBrowser:
|
||
def __init__(self, page):
|
||
self.page = page
|
||
self.closed = False
|
||
|
||
def new_page(self, viewport=None):
|
||
return self.page
|
||
|
||
def close(self):
|
||
self.closed = True
|
||
|
||
|
||
class _FakePW:
|
||
def __init__(self, page, fail_launch=False):
|
||
self.page = page
|
||
self.fail_launch = fail_launch
|
||
self.stopped = False
|
||
self.launch_kwargs = None
|
||
|
||
class _Chromium:
|
||
def __init__(self, outer):
|
||
self.outer = outer
|
||
|
||
def launch(self, **kw):
|
||
self.outer.launch_kwargs = kw
|
||
if self.outer.fail_launch:
|
||
raise RuntimeError("no chromium")
|
||
return _FakeBrowser(self.outer.page)
|
||
|
||
self.chromium = _Chromium(self)
|
||
|
||
def start(self):
|
||
return self
|
||
|
||
def stop(self):
|
||
self.stopped = True
|
||
|
||
|
||
class PlaywrightBrowserTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.page = _FakePage()
|
||
self.env = mock.patch.dict(os.environ, {"HTTPS_PROXY": "",
|
||
"HTTP_PROXY": ""})
|
||
self.env.start()
|
||
|
||
def tearDown(self):
|
||
self.env.stop()
|
||
self.tmp.cleanup()
|
||
|
||
def _browser(self, proxy=None, fail_launch=False):
|
||
return pw_mod.Browser(
|
||
cwd=self.tmp.name, proxy=proxy,
|
||
browser_factory=lambda: _FakePW(self.page, fail_launch))
|
||
|
||
def test_navigate_and_info(self):
|
||
b = self._browser(proxy="http://proxy:1")
|
||
out = b.navigate("https://example.com")
|
||
self.assertTrue(out["ok"])
|
||
self.assertEqual(out["url"], "https://example.com")
|
||
self.assertEqual(out["title"], "Fake Title")
|
||
self.assertEqual(b.cached_info()["url"], "https://example.com")
|
||
self.assertTrue(b.cached_info()["started"])
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.navigate(" ")
|
||
self.page.goto_fail = True
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.navigate("https://example.com")
|
||
|
||
def test_launch_failure(self):
|
||
b = self._browser(fail_launch=True)
|
||
with self.assertRaises(pw_mod.BrowserError) as cm:
|
||
b.snapshot()
|
||
self.assertIn("browser launch failed", str(cm.exception))
|
||
|
||
def test_snapshot_truncation_and_failure(self):
|
||
b = self._browser()
|
||
out = b.snapshot()
|
||
self.assertEqual(out["text"], "hello world")
|
||
self.assertFalse(out["truncated"])
|
||
self.page.body_text = "x" * 9000
|
||
out = b.snapshot()
|
||
self.assertTrue(out["truncated"])
|
||
self.assertEqual(len(out["text"]), 8000)
|
||
orig = self.page.locator
|
||
self.page.locator = lambda s: (_ for _ in ()).throw(
|
||
RuntimeError("loc fail"))
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.snapshot()
|
||
self.page.locator = orig
|
||
|
||
def test_click_fill_key_type(self):
|
||
b = self._browser()
|
||
self.assertTrue(b.click("button.ok")["ok"])
|
||
self.assertIn("button.ok", self.page.locator_clicks)
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.click("")
|
||
self.assertTrue(b.fill("#q", None)["ok"])
|
||
self.assertEqual(self.page.filled, ("#q", ""))
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.fill("", "x")
|
||
self.assertTrue(b.key_press("Enter")["ok"])
|
||
self.assertEqual(self.page.keyboard.pressed, ["Enter"])
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.key_press(" ")
|
||
self.assertTrue(b.type_text("hi")["ok"])
|
||
self.assertEqual(self.page.keyboard.typed, ["hi"])
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.type_text(None)
|
||
|
||
def test_mouse_and_click_at(self):
|
||
b = self._browser()
|
||
out = b.mouse_move(10, 20)
|
||
self.assertEqual((out["x"], out["y"]), (10.0, 20.0))
|
||
self.assertEqual(self.page.mouse.moves, [(10.0, 20.0)])
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.mouse_move("a", 1)
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.mouse_move(2000, 10)
|
||
out = b.click_at(5, 6)
|
||
self.assertTrue(out["ok"])
|
||
self.assertEqual(self.page.mouse.clicks, [(5.0, 6.0)])
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.click_at("x", 1)
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.click_at(10, 900)
|
||
|
||
def test_snapshot_refs_and_click_ref(self):
|
||
b = self._browser()
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.click_ref("e0")
|
||
self.page.eval_result = [
|
||
{"tag": "button", "role": "", "name": "OK",
|
||
"path": "body > button"},
|
||
{"tag": "a", "role": "link", "name": "Home",
|
||
"path": "body > a"}]
|
||
out = b.snapshot_refs()
|
||
self.assertEqual(out["count"], 2)
|
||
self.assertEqual(out["elements"][0]["ref"], "e0")
|
||
out = b.click_ref("e1")
|
||
self.assertTrue(out["ok"])
|
||
self.assertEqual(out["ref"], "e1")
|
||
self.assertIn("body > a", self.page.locator_clicks)
|
||
|
||
def test_screencast_sink_and_pump(self):
|
||
b = self._browser()
|
||
self.assertEqual(b.pump_screencast(), "")
|
||
sink = _SimpleQueue()
|
||
b.set_frame_sink(sink)
|
||
b.navigate("https://example.com")
|
||
# navigate 后 sink 已接上 CDP（_start_screencast 在
|
||
# navigate 内跑）；再显式 set 一次走“已接同一页”短路。
|
||
b.set_frame_sink(sink)
|
||
cdp = self.page.context.cdp
|
||
self.assertIsNotNone(cdp)
|
||
cdp.handlers["Page.screencastFrame"](
|
||
{"sessionId": "s1", "data": "frame-b64"})
|
||
self.assertEqual(b.pump_screencast(), "frame-b64")
|
||
self.assertIn(("Page.screencastFrameAck",
|
||
{"sessionId": "s1"}), cdp.sent)
|
||
# eager 帧（空 session）只存不 ack。
|
||
sink.put(("", "eager-b64"))
|
||
self.assertEqual(b.pump_screencast(), "eager-b64")
|
||
self.assertEqual(
|
||
[s for s in cdp.sent if s[0] == "Page.screencastFrameAck"],
|
||
[("Page.screencastFrameAck", {"sessionId": "s1"})])
|
||
|
||
def test_live_shot_screenshot_evaluate_close(self):
|
||
b = self._browser()
|
||
status = b.live_shot()
|
||
self.assertEqual(status["image"], "browser-live.png")
|
||
self.assertEqual(status["bytes"], 7)
|
||
with open(status["status_path"]) as f:
|
||
sidecar = json.load(f)
|
||
self.assertEqual(sidecar["image"], "browser-live.png")
|
||
shot = b.screenshot()
|
||
self.assertTrue(shot["filename"].startswith("shot-"))
|
||
self.assertEqual(shot["bytes"], 7)
|
||
shot2 = b.screenshot("my.png")
|
||
self.assertEqual(shot2["filename"], "my.png")
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.screenshot("../escape.png")
|
||
self.assertEqual(b.evaluate("1+1")["result"], "eval-ok")
|
||
with self.assertRaises(pw_mod.BrowserError):
|
||
b.evaluate(" ")
|
||
out = b.close()
|
||
self.assertTrue(out["closed"])
|
||
self.assertEqual(out["errors"], [])
|
||
# 已关闭后再次 close 仍返回 closed（各 closer 遇 None 短路）。
|
||
out = b.close()
|
||
self.assertTrue(out["closed"])
|
||
|
||
|
||
class PlaywrightEnsureTest(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.src = os.path.join(self.tmp.name, "src")
|
||
self.home = os.path.join(self.tmp.name, "home")
|
||
os.makedirs(self.src)
|
||
for name in ("playwright_browser.py", "playwright_mcp.py",
|
||
"playwright_remote.py"):
|
||
with open(os.path.join(self.src, name), "w") as f:
|
||
f.write("# %s v1\n" % name)
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def test_seed_refresh_and_config(self):
|
||
out = pw_mod.ensure_codex_home_playwright(
|
||
self.home, script_dir=self.src)
|
||
self.assertTrue(out["script_changed"])
|
||
self.assertTrue(out["config_changed"])
|
||
with open(os.path.join(self.home, "config.toml")) as f:
|
||
cfg = f.read()
|
||
self.assertIn("[mcp_servers.playwright]", cfg)
|
||
# 第二次：脚本相同、配置已有 section，全部不变。
|
||
out = pw_mod.ensure_codex_home_playwright(
|
||
self.home, script_dir=self.src)
|
||
self.assertFalse(out["script_changed"])
|
||
self.assertFalse(out["config_changed"])
|
||
# 旧脚本被刷新上报。
|
||
with open(os.path.join(self.home,
|
||
"playwright_mcp.py"), "w") as f:
|
||
f.write("# stale\n")
|
||
out = pw_mod.ensure_codex_home_playwright(
|
||
self.home, script_dir=self.src)
|
||
self.assertIn("playwright_mcp.py", out["script_refreshed"])
|
||
# 自定义 command_args 写进 section（删旧配置重写，
|
||
# 否则已有 section 不会被改写）。
|
||
os.remove(os.path.join(self.home, "config.toml"))
|
||
out = pw_mod.ensure_codex_home_playwright(
|
||
self.home, command_args=["py3", "mcp.py", "--x"],
|
||
script_dir=self.src)
|
||
self.assertTrue(out["config_changed"])
|
||
with open(os.path.join(self.home, "config.toml")) as f:
|
||
cfg = f.read()
|
||
self.assertIn('"py3"', cfg)
|
||
|
||
def test_config_no_trailing_newline(self):
|
||
os.makedirs(self.home)
|
||
with open(os.path.join(self.home, "config.toml"), "w") as f:
|
||
f.write("[other]")
|
||
pw_mod.ensure_codex_home_playwright(self.home,
|
||
script_dir=self.src)
|
||
with open(os.path.join(self.home, "config.toml")) as f:
|
||
cfg = f.read()
|
||
self.assertIn("[other]\n\n[mcp_servers.playwright]", cfg)
|
||
|
||
def test_missing_src_script_skipped(self):
|
||
os.remove(os.path.join(self.src, "playwright_remote.py"))
|
||
out = pw_mod.ensure_codex_home_playwright(
|
||
self.home, script_dir=self.src)
|
||
self.assertTrue(out["script_changed"])
|
||
self.assertFalse(os.path.exists(
|
||
os.path.join(self.home, "playwright_remote.py")))
|
||
|
||
def test_mcp_env_writer(self):
|
||
os.makedirs(self.home)
|
||
self.assertFalse(pw_mod.ensure_playwright_mcp_env(
|
||
self.home, {}))
|
||
self.assertTrue(pw_mod.ensure_playwright_mcp_env(
|
||
self.home, {"B": "2", "A": "1"}))
|
||
with open(os.path.join(self.home, "config.toml")) as f:
|
||
cfg = f.read()
|
||
self.assertIn("[mcp_servers.playwright.env]", cfg)
|
||
self.assertLess(cfg.index('A = '), cfg.index('B = '))
|
||
# 幂等替换：旧 env 段被新段取代，不残留旧 key。
|
||
self.assertTrue(pw_mod.ensure_playwright_mcp_env(
|
||
self.home, {"C": "3"}))
|
||
with open(os.path.join(self.home, "config.toml")) as f:
|
||
cfg = f.read()
|
||
self.assertIn('C = ', cfg)
|
||
self.assertNotIn('A = ', cfg)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|