Files
qmt_bridge/tests/test_ws_real.py

95 lines
3.0 KiB
Python

# -*- coding: utf-8 -*-
"""Real QMT WS push test: connect /ws, subscribe whole, receive pushes."""
import socket, struct, base64, os, json, time, urllib.request
def ws_connect(port):
s = socket.create_connection(("127.0.0.1", port), timeout=5)
key = base64.b64encode(os.urandom(16)).decode()
req = ("GET /ws HTTP/1.1\r\nHost: 127.0.0.1:%d\r\nUpgrade: websocket\r\n"
"Connection: Upgrade\r\nSec-WebSocket-Key: %s\r\n"
"Sec-WebSocket-Version: 13\r\n\r\n" % (port, key))
s.sendall(req.encode())
resp = s.recv(4096)
assert b"101" in resp, resp[:100]
return s
def ws_read_frame(s, timeout=15):
"""Read one frame; returns (opcode, payload_str) or None."""
s.settimeout(timeout)
hdr = s.recv(2)
if len(hdr) < 2:
return None
opcode = hdr[0] & 0x0F
b1 = hdr[1]
ln = b1 & 0x7F
if ln == 126:
ln = struct.unpack(">H", s.recv(2))[0]
elif ln == 127:
ln = struct.unpack(">Q", s.recv(8))[0]
payload = b""
while len(payload) < ln:
chunk = s.recv(ln - len(payload))
if not chunk:
break
payload += chunk
return (opcode, payload.decode("utf-8", "ignore"))
# 1. connect
ws = ws_connect(8610)
print("WS connected")
# 2. subscribe whole
req = urllib.request.Request(
"http://127.0.0.1:8610/data/subscribe",
data=json.dumps({"type": "whole", "codes": ["SH", "SZ"]}).encode(),
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=10) as r:
resp = json.loads(r.read().decode())
print("subscribe:", resp)
sub_id = resp.get("sub_id")
# 3. receive frames for ~20s, looking for whole pushes
print("waiting for whole pushes (up to 20s)...")
got = 0
deadline = time.time() + 20
while time.time() < deadline and got < 3:
try:
frame = ws_read_frame(ws, timeout=5)
except socket.timeout:
continue
if frame is None:
break
opcode, payload = frame
if opcode == 0x1: # text
try:
parsed = json.loads(payload)
t = parsed.get("type")
if t == "whole" and parsed.get("data"):
codes = list(parsed["data"].keys())
sample = parsed["data"].get(codes[0], {}) if codes else {}
print("whole push: %d codes, sample=%s" % (len(codes), json.dumps(sample, ensure_ascii=False)[:120]))
got += 1
else:
print("text frame type=%s" % t)
except Exception as e:
print("parse err:", e, "payload:", payload[:80])
elif opcode == 0x9: # ping -> pong
ws.sendall(bytes([0x8A, 0x00]))
print("ping -> pong")
elif opcode == 0x8:
print("close frame")
break
# 4. unsubscribe
req = urllib.request.Request(
"http://127.0.0.1:8610/data/unsubscribe",
data=json.dumps({"sub_id": sub_id}).encode(),
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=10) as r:
print("unsubscribe:", r.read().decode())
ws.close()
print("=== DONE got_whole=%d ===" % got)