xianyan/docs/toolsapi/scripts/test_full_e2e.py

"""
闲言APP - 文件传输全流程E2E测试脚本
创建时间: 2026-05-19
作用: 模拟两个设备从发现→配对→消息→文件→画布→屏幕共享完整流程
"""
import asyncio
import json
import time

WS_URL = 'wss://tools.wktyl.com:9443'

async def ws_connect():
    import websockets
    return await websockets.connect(WS_URL, ping_interval=None)

async def register_device(ws, device_id, alias, fingerprint, user_id=None, ip_city=None, ip_range=None):
    msg = {
        'type': 'register',
        'data': {'deviceId': device_id, 'alias': alias, 'deviceType': 'mobile', 'fingerprint': fingerprint, 'protocol': 'xianyan-v1'}
    }
    if user_id:
        msg['data']['userId'] = user_id
    if ip_city:
        msg['data']['ipCity'] = ip_city
    if ip_range:
        msg['data']['ipRange'] = ip_range
    await ws.send(json.dumps(msg))
    for _ in range(5):
        resp = json.loads(await asyncio.wait_for(ws.recv(), timeout=10))
        if resp.get('type') == 'registered':
            return resp
    return None

async def drain_messages(ws, count=3, timeout=5):
    messages = []
    for _ in range(count):
        try:
            msg = json.loads(await asyncio.wait_for(ws.recv(), timeout=timeout))
            messages.append(msg)
        except asyncio.TimeoutError:
            break
    return messages

async def test_full_e2e():
    """全流程: 注册→配对码配对→文本消息→文件元数据→画布同步→屏幕共享"""
    ts = int(time.time())

    ws_a = await ws_connect()
    ws_b = await ws_connect()

    print('  [1/7] 注册设备...')
    reg_a = await register_device(ws_a, f'e2e-a-{ts}', 'iPhone 16 Pro', f'fp_e2e_a_{ts}', user_id=f'user-e2e-{ts}', ip_city='北京', ip_range='110.123.45.0/24')
    reg_b = await register_device(ws_b, f'e2e-b-{ts}', 'MacBook Pro', f'fp_e2e_b_{ts}', user_id=f'user-e2e-{ts}', ip_city='北京', ip_range='110.123.45.0/24')
    assert reg_a is not None and reg_b is not None, 'Registration failed'
    print('  [PASS] 设备注册成功')

    print('  [2/7] 配对码配对...')
    await ws_a.send(json.dumps({'type': 'pairing-code-create', 'data': {'alias': 'iPhone 16 Pro'}}))
    code = None
    for _ in range(5):
        msg = json.loads(await asyncio.wait_for(ws_a.recv(), timeout=10))
        if msg.get('type') == 'pairing-code-created':
            code = msg['pairingCode']
            break
    assert code is not None, 'Failed to create pairing code'

    await ws_b.send(json.dumps({'type': 'pairing-code-join', 'data': {'pairingCode': code}}))
    matched = False
    for _ in range(5):
        msg_b = json.loads(await asyncio.wait_for(ws_b.recv(), timeout=10))
        if msg_b.get('type') == 'pairing-matched' and msg_b.get('success'):
            matched = True
            break
    for _ in range(5):
        msg_a = json.loads(await asyncio.wait_for(ws_a.recv(), timeout=10))
        if msg_a.get('type') == 'pairing-matched' and msg_a.get('success'):
            break
    assert matched, 'Pairing code match failed'
    print(f'  [PASS] 配对码配对成功: {code}')

    print('  [3/7] 文本消息...')
    await ws_a.send(json.dumps({
        'type': 'textMessage',
        'to': f'fp_e2e_b_{ts}',
        'message': {'text': 'Hello from iPhone!'},
        'timestamp': int(time.time() * 1000)
    }))
    received = False
    for _ in range(5):
        msg = json.loads(await asyncio.wait_for(ws_b.recv(), timeout=10))
        if msg.get('type') == 'textMessage':
            received = True
            assert msg.get('from') == f'fp_e2e_a_{ts}', f'Wrong from: {msg}'
            break
    assert received, 'Text message not received'
    print('  [PASS] 文本消息收发成功')

    print('  [4/7] 文件元数据...')
    await ws_a.send(json.dumps({
        'type': 'fileMeta',
        'to': f'fp_e2e_b_{ts}',
        'file': {'fileName': 'photo.jpg', 'fileSize': 1024000, 'mimeType': 'image/jpeg'},
        'timestamp': int(time.time() * 1000)
    }))
    received = False
    for _ in range(5):
        msg = json.loads(await asyncio.wait_for(ws_b.recv(), timeout=10))
        if msg.get('type') == 'fileMeta':
            received = True
            break
    assert received, 'File meta not received'
    print('  [PASS] 文件元数据传输成功')

    print('  [5/7] 画布同步...')
    canvas_id = f'e2e-canvas-{ts}'
    await ws_a.send(json.dumps({'type': 'canvas-join', 'payload': {'canvasId': canvas_id}}))
    await drain_messages(ws_a, 3)
    await ws_b.send(json.dumps({'type': 'canvas-join', 'payload': {'canvasId': canvas_id}}))
    await drain_messages(ws_a, 3)
    await drain_messages(ws_b, 3)

    await ws_a.send(json.dumps({
        'type': 'canvas-stroke',
        'payload': {'canvasId': canvas_id, 'stroke': {'strokeId': 'e2e-s1', 'points': [[50, 60]], 'color': '#000', 'width': 2, 'tool': 'pen'}}
    }))
    received = False
    for _ in range(5):
        msg = json.loads(await asyncio.wait_for(ws_b.recv(), timeout=10))
        if msg.get('type') == 'canvas-stroke':
            received = True
            break
    assert received, 'Canvas stroke not received'
    print('  [PASS] 画布同步成功')

    print('  [6/7] 屏幕共享请求...')
    await ws_a.send(json.dumps({
        'type': 'screen-share-request',
        'to': f'fp_e2e_b_{ts}',
        'payload': {'direction': 'view'}
    }))
    received = False
    for _ in range(5):
        msg = json.loads(await asyncio.wait_for(ws_b.recv(), timeout=10))
        if msg.get('type') == 'screen-share-request':
            received = True
            break
    assert received, 'Screen share request not received'

    await ws_b.send(json.dumps({
        'type': 'screen-share-accept',
        'to': f'fp_e2e_a_{ts}',
        'payload': {'direction': 'view'}
    }))
    received = False
    for _ in range(5):
        msg = json.loads(await asyncio.wait_for(ws_a.recv(), timeout=10))
        if msg.get('type') == 'screen-share-accept':
            received = True
            break
    assert received, 'Screen share accept not received'
    print('  [PASS] 屏幕共享请求/接受成功')

    print('  [7/7] 心跳保活...')
    await ws_a.send(json.dumps({'type': 'heartbeat', 'data': {'timestamp': int(time.time() * 1000)}}))
    got_ack = False
    for _ in range(5):
        msg = json.loads(await asyncio.wait_for(ws_a.recv(), timeout=10))
        if msg.get('type') == 'heartbeat_ack':
            got_ack = True
            break
    assert got_ack, 'Heartbeat ack not received'
    print('  [PASS] 心跳保活成功')

    await ws_a.close()
    await ws_b.close()
    print('\n  ===== 全流程E2E测试通过! =====')

async def run_all_tests():
    print('\n===== 全流程E2E测试 =====')
    await test_full_e2e()

if __name__ == '__main__':
    asyncio.run(run_all_tests())