ykxmz/_build/test_file_pull.py

# -*- coding: utf-8 -*-
"""正确地复现 file_pull：上传 → 只读一条回复 → 拉取 → 收分块事件。

上一版我多发了一条空命令，把回复错位了，测出来的结果是假的。
这次每步只读该读的。
"""
from __future__ import annotations

import base64
import json
import shutil
import socket
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(r"C:\Users\shu\Desktop\deepseekagent\LanRemote")
EXE = ROOT / "dist_exe" / "ykxmz" / "ykxmz.exe"
PORT = 8817


class Talk:
    def __init__(self, port):
        self.sock = socket.create_connection(("127.0.0.1", port), timeout=20)
        self.sock.settimeout(25)
        self.buf = b""

    def read_one(self, timeout=25):
        deadline = time.time() + timeout
        while True:
            while b"\n" not in self.buf:
                if time.time() > deadline:
                    raise TimeoutError("超时")
                chunk = self.sock.recv(262144)
                if not chunk:
                    raise ConnectionError("断开")
                self.buf += chunk
            line, _, self.buf = self.buf.partition(b"\n")
            return json.loads(line.decode("utf-8"))

    def call(self, obj, timeout=25):
        self.sock.sendall((json.dumps(obj, ensure_ascii=False) + "\n").encode("utf-8"))
        deadline = time.time() + timeout
        while True:
            message = self.read_one(timeout=max(0.5, deadline - time.time()))
            if "event" in message:
                continue
            return message


def main() -> int:
    d = EXE.parent
    real = d / "data"
    recv = d / "received"
    bak = None
    if real.exists():
        bak = d / "data_bk_fp2"
        if bak.exists():
            shutil.rmtree(bak, ignore_errors=True)
        shutil.move(str(real), str(bak))
    if recv.exists():
        shutil.rmtree(recv, ignore_errors=True)

    proc = None
    try:
        proc = subprocess.Popen(
            [str(EXE), "--port", str(PORT), "--add-device", "10"],
            cwd=str(d), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        code = ""
        for _ in range(40):
            time.sleep(0.5)
            try:
                code = json.loads((d / "data" / "settings.json").read_text(
                    encoding="utf-8")).get("pairing_code", "")
            except Exception:                             # noqa: BLE001
                code = ""
            if code:
                try:
                    probe = socket.create_connection(("127.0.0.1", PORT), timeout=1)
                    probe.close()
                    break
                except OSError:
                    continue

        talk = Talk(PORT)
        talk.call({"cmd": "pair", "deviceId": "fp", "deviceName": "fp", "code": code})

        payload = bytes((i * 31 + 7) & 0xFF for i in range(300000))
        print(f"  ① 上传 {len(payload)} 字节")
        chunk = 256 * 1024
        total = max(1, (len(payload) + chunk - 1) // chunk)
        for index in range(total):
            piece = payload[index * chunk:(index + 1) * chunk]
            talk.sock.sendall((json.dumps({
                "cmd": "file_push", "name": "java_payload.bin",
                "seq": index, "last": index == total - 1,
                "data": base64.b64encode(piece).decode()}) + "\n").encode("utf-8"))
        reply = talk.call({}, timeout=20) if False else talk.read_one(timeout=20)
        # 最后一块的回复
        print(f"     上传回复: {str(reply)[:90]}")

        print()
        print("  ② file_pull 拉回来")
        talk.sock.sendall((json.dumps({
            "cmd": "file_pull", "path": "java_payload.bin"}) + "\n").encode("utf-8"))

        events = []
        got = bytearray()
        done = False
        final = None
        deadline = time.time() + 25
        while time.time() < deadline:
            message = talk.read_one(timeout=max(0.5, deadline - time.time()))
            if "event" in message:
                if message.get("event") == "file":
                    got += base64.b64decode(message.get("data", ""))
                    events.append((message.get("seq"), message.get("last")))
                    if message.get("last"):
                        done = True
                        break
                continue
            final = message
            break

        print(f"      收到分块 {len(events)} 个: {events[:4]}")
        print(f"      收齐标记 last: {done}")
        print(f"      字节数: {len(got)}（期望 {len(payload)}）")
        print(f"      内容一致: {bytes(got) == payload}")
        print(f"      最终回复: {str(final)[:90]}")
        talk.sock.close()

        ok = done and bytes(got) == payload
        print()
        print(f"  >>> file_pull {'正常' if ok else '有问题'}")
        return 0 if ok else 1
    finally:
        if proc:
            try:
                proc.terminate()
                proc.wait(timeout=8)
            except Exception:                             # noqa: BLE001
                try:
                    proc.kill()
                except Exception:                     # noqa: BLE001
                    pass
        time.sleep(1)
        if real.exists():
            shutil.rmtree(real, ignore_errors=True)
        if recv.exists():
            shutil.rmtree(recv, ignore_errors=True)
        if bak and bak.exists():
            shutil.move(str(bak), str(real))


if __name__ == "__main__":
    sys.exit(main())