init

tests/repl_drive.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""repl_drive.py - 用 pty 驱动 ligent 的 REPL，跑多轮对话并打印全部输出。
+
+交互式行为（多轮上下文、/命令、行编辑）没法用管道测 —— 管道模式会把
+整段 stdin 当成一个任务。所以这里开一个伪终端，逐条发送提示词。
+
+用法：
+    repl_drive.py --bin ./ligent --base-url URL --prompt "第一轮" --prompt "第二轮"
+"""
+import argparse
+import os
+import pty
+import re
+import select
+import sys
+import time
+
+
+def drive(binary, args, prompts, wait_per_prompt=2.0, quiet_after=1.5,
+          raw_keys=None):
+    pid, fd = pty.fork()
+    if pid == 0:
+        os.execv(binary, [binary] + args)
+        os._exit(127)
+
+    out = bytearray()
+    # 行编辑器每次回到提示符都会打开括号粘贴，所以这个序列就是"轮到用户输入了"
+    # 的信号。用它做同步点，比"等输出安静下来"可靠也快得多。
+    PROMPT_READY = b"\x1b[?2004h"
+    ready_seen = 0
+    tail = b""
+
+    def note(chunk):
+        # 标记可能被 read 切断，所以要连着上一块的尾巴一起找
+        nonlocal tail, ready_seen
+        buf = tail + chunk
+        ready_seen += buf.count(PROMPT_READY)
+        tail = buf[-(len(PROMPT_READY) - 1):]
+
+    def drain(timeout, idle=0.5, until_ready=False):
+        """等到"回到提示符"（until_ready），或者输出彻底安静下来。
+
+        必须按标记等，不能只按安静等：等待动画每 250ms 画一帧，pty 永远
+        不会安静超过 idle，于是每次都要白等满 timeout —— 而 mock 其实几
+        毫秒就答复了。这是整个 e2e 从 126 秒掉下来的关键。
+        """
+        want = ready_seen + 1 if until_ready else None
+        deadline = time.time() + timeout
+        last = time.time()
+        while time.time() < deadline:
+            r, _, _ = select.select([fd], [], [], 0.02)
+            if r:
+                try:
+                    chunk = os.read(fd, 65536)
+                except OSError:
+                    return False
+                if not chunk:
+                    return False
+                out.extend(chunk)
+                note(chunk)
+                last = time.time()
+                if want is not None and ready_seen >= want:
+                    return True
+            elif time.time() - last > idle:
+                return True
+        return True
+
+    drain(quiet_after, until_ready=True)    # 等横幅画完、第一次回到提示符
+    if raw_keys is not None:
+        # 逐字节发送原始按键序列（用于测方向键等编辑行为）
+        for chunk, pause in raw_keys:
+            os.write(fd, chunk if isinstance(chunk, bytes) else chunk.encode("utf-8"))
+            drain(pause)
+        os.write(fd, b"/exit\n")
+        drain(0.8)
+        try:
+            os.close(fd)
+        except OSError:
+            pass
+        try:
+            os.waitpid(pid, 0)
+        except ChildProcessError:
+            pass
+        return re.sub(r"\x1b\[[0-9;]*m", "", out.decode("utf-8", "replace"))
+
+    for p in prompts:
+        os.write(fd, (p + "\n").encode("utf-8"))
+        if not drain(wait_per_prompt, until_ready=True):
+            break
+    os.write(fd, b"/exit\n")
+    drain(0.8)
+
+    try:
+        os.close(fd)
+    except OSError:
+        pass
+    try:
+        os.waitpid(pid, 0)
+    except ChildProcessError:
+        pass
+
+    text = out.decode("utf-8", "replace")
+    return re.sub(r"\x1b\[[0-9;]*m", "", text)      # 去掉 ANSI 样式
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--bin", default="./ligent")
+    ap.add_argument("--base-url", required=True)
+    ap.add_argument("--api-key", default="test")
+    ap.add_argument("--prompt", action="append", default=[])
+    ap.add_argument("--extra", action="append", default=[])
+    ap.add_argument("--wait", type=float, default=2.0)
+    ap.add_argument("--keys", action="append", default=[],
+                    help="原始按键序列（支持 \\x1b 转义），按顺序发送")
+    a = ap.parse_args()
+
+    # --no-history：被驱动的实例不该往用户真实的 ligent.history 里写测试垃圾。
+    # 需要历史的调用方用 --extra=--history PATH（排在后面，会覆盖这一条）。
+    args = ["--api-key", a.api_key, "--base-url", a.base_url, "--no-color",
+            "--no-confirm", "--no-auto-save", "--no-history"] + a.extra
+    if a.keys:
+        raw = [(k.encode("utf-8").decode("unicode_escape").encode("latin-1"), 0.25)
+               for k in a.keys]
+        sys.stdout.write(drive(a.bin, args, [], a.wait, raw_keys=raw))
+    else:
+        sys.stdout.write(drive(a.bin, args, a.prompt, a.wait))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())