kyugao/qmt_bridge
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: init qmt_bridge repo (HTTP+WS bridge, MCP endpoint, docs, references)

Browse Source
This commit is contained in:
Docker
2026-08-26 16:53:15 +08:00
commit 22a5b8ca04
210 changed files with 68176 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_mcp.py
+209
View File
@@ -0,0 +1,209 @@
# -*- coding: utf-8 -*-
"""Local mock test for the MCP endpoint (POST /mcp JSON-RPC).
Covers: initialize, tools/list, tools/call for all 9 tools, plus
error paths (parse error, unknown method, missing required code) and
notification 202. Run with local python (3.6-compatible stdlib only).
"""
import sys, io, os, json, time, urllib.request
if sys.stdout is None:
sys.stdout = io.StringIO()
if sys.stderr is None:
sys.stderr = io.StringIO()
SRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src")
sys.path.insert(0, SRC)
import bridge_util, bridge_http_server, bridge_data_adapter
class FakeDF:
"""Minimal pandas.DataFrame stand-in (columns/shape/iat/index) so the
test runs on stdlib-only Python (no pandas needed locally)."""
def __init__(self, columns, rows):
self.columns = columns
self._rows = rows
@property
def shape(self):
return (len(self._rows), len(self.columns))
def __len__(self):
return len(self._rows)
def __iter__(self):
return iter(self.columns)
@property
def index(self):
return ["20260818", "20260819"]
def iat(self, i, j):
return self._rows[i][j]
class FakeCtx:
def get_market_data_ex(self, fields, codes, **kw):
cols = ["open", "high", "low", "close", "volume", "amount"]
rows = [[10.0, 10.5, 9.8, 10.2, 100.0, 1000.0],
[11.0, 11.5, 10.6, 11.2, 200.0, 2200.0]]
return {codes[0]: FakeDF(cols, rows)}
def get_full_tick(self, codes):
return {c: {"lastPrice": 10.2, "volume": 500,
"askPrice": [10.3], "bidPrice": [10.1]} for c in codes}
def get_instrument_detail(self, code):
return {"InstrumentName": "Test", "UpStopPrice": 11.0,
"DownStopPrice": 9.0, "OpenDate": "20010827"}
def get_trading_dates(self, *a, **kw):
return ["20260818", "20260819", "20260820"]
bridge_util.CTX = FakeCtx()
bridge_util.ACCOUNT = "TESTACC"
bridge_util.ACCOUNT_TYPE = "STOCK"
def fake_trade(acct, atype, dtype, strategy_name=None):
if dtype == "position":
return [type("Pos", (), {"m_strInstrumentID": "600519", "m_strInstrumentName": "Moutai",
"m_nVolume": 100, "m_nCanUseVolume": 50,
"m_nFrozenVolume": 50, "m_dOpenPrice": 1500.0,
"m_dFloatProfit": 1000.0})()]
if dtype == "order":
return [type("O", (), {"m_strInstrumentID": "600519", "m_nOrderStatus": 48})()]
return [type("D", (), {"m_strInstrumentID": "600519", "m_nVolume": 100,
"m_dPrice": 1500.0})()]
bridge_util.QMT_API["get_trade_detail_data"] = fake_trade
bridge_http_server.PORT = 18692
bridge_http_server.TOKEN = ""
bridge_http_server.ACCOUNT = "TESTACC"
bridge_util._refresh_trade_cache("TESTACC", "STOCK")
bridge_http_server.start_server()
time.sleep(0.8)
BASE = "http://127.0.0.1:18692/mcp"
def post(payload):
data = json.dumps(payload).encode()
req = urllib.request.Request(BASE, data=data,
headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=5) as r:
return r.status, json.loads(r.read().decode())
def check(name, cond, detail=""):
if not cond:
raise AssertionError("%s FAILED %s" % (name, detail))
print("OK %s" % name)
# 1. initialize
st, r = post({"jsonrpc": "2.0", "id": 1, "method": "initialize",
"params": {"protocolVersion": "2025-06-18",
"capabilities": {}, "clientInfo": {"name": "t", "version": "1"}}})
check("initialize", st == 200 and r["result"]["protocolVersion"] == "2025-06-18"
and r["result"]["capabilities"].get("tools") is not None,
str(r))
# 2. tools/list -> 9 tools
st, r = post({"jsonrpc": "2.0", "id": 2, "method": "tools/list"})
tools = r["result"]["tools"]
names = [t["name"] for t in tools]
check("tools/list count", len(tools) == 9, str(names))
for need in ["qmt_kline", "qmt_quote", "qmt_tick", "qmt_instrument",
"qmt_trading_dates", "qmt_positions", "qmt_asset",
"qmt_orders", "qmt_trades"]:
check("tool %s" % need, need in names)
# 3. tools/call qmt_kline
st, r = post({"jsonrpc": "2.0", "id": 3, "method": "tools/call",
"params": {"name": "qmt_kline",
"arguments": {"code": "600519", "period": "1d", "count": 5}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_kline", text["code"] == "600519.SH" and text["count"] == 2, str(text)[:120])
# 4. tools/call qmt_quote
st, r = post({"jsonrpc": "2.0", "id": 4, "method": "tools/call",
"params": {"name": "qmt_quote", "arguments": {"code": "600519.SH"}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_quote", text["ok"] is True and text["data"]["lastPrice"] == 10.2, str(text)[:120])
# 5. tools/call qmt_tick
st, r = post({"jsonrpc": "2.0", "id": 5, "method": "tools/call",
"params": {"name": "qmt_tick",
"arguments": {"codes": "600000.SH,000001.SZ"}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_tick", text["ok"] is True and len(text["data"]) == 2, str(text)[:120])
# 6. tools/call qmt_instrument
st, r = post({"jsonrpc": "2.0", "id": 6, "method": "tools/call",
"params": {"name": "qmt_instrument", "arguments": {"code": "600519"}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_instrument", text["data"]["UpStopPrice"] == 11.0, str(text)[:120])
# 7. tools/call qmt_trading_dates
st, r = post({"jsonrpc": "2.0", "id": 7, "method": "tools/call",
"params": {"name": "qmt_trading_dates",
"arguments": {"start": "20260818", "end": "20260820"}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_trading_dates", text["count"] == 3, str(text)[:120])
# 8. tools/call qmt_positions
st, r = post({"jsonrpc": "2.0", "id": 8, "method": "tools/call",
"params": {"name": "qmt_positions", "arguments": {}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_positions", text["data"]["summary"]["count"] == 1, str(text)[:120])
# 9. tools/call qmt_asset
st, r = post({"jsonrpc": "2.0", "id": 9, "method": "tools/call",
"params": {"name": "qmt_asset", "arguments": {}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_asset", "data" in text, str(text)[:120])
# 10. tools/call qmt_orders (with code filter)
st, r = post({"jsonrpc": "2.0", "id": 10, "method": "tools/call",
"params": {"name": "qmt_orders",
"arguments": {"code": "600519", "status": "active"}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_orders", text["ok"] is True and len(text["data"]) == 1, str(text)[:120])
# 11. tools/call qmt_trades
st, r = post({"jsonrpc": "2.0", "id": 11, "method": "tools/call",
"params": {"name": "qmt_trades", "arguments": {}}})
text = json.loads(r["result"]["content"][0]["text"])
check("qmt_trades", text["ok"] is True, str(text)[:120])
# 12. error: parse error (bad json)
data = b"{not json"
req = urllib.request.Request(BASE, data=data,
headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=5) as rr:
err = json.loads(rr.read().decode())
check("parse error", err["error"]["code"] == -32700, str(err))
# 13. error: unknown method
st, r = post({"jsonrpc": "2.0", "id": 13, "method": "bogus/method"})
check("unknown method", r["error"]["code"] == -32601, str(r))
# 14. error: missing required code
st, r = post({"jsonrpc": "2.0", "id": 14, "method": "tools/call",
"params": {"name": "qmt_kline", "arguments": {}}})
check("missing code", r["error"]["code"] == -32602, str(r))
# 15. error: unknown tool
st, r = post({"jsonrpc": "2.0", "id": 15, "method": "tools/call",
"params": {"name": "qmt_nope", "arguments": {}}})
check("unknown tool", r["error"]["code"] == -32602, str(r))
# 16. notification -> 202 no body
data = json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized"}).encode()
req = urllib.request.Request(BASE, data=data,
headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=5) as rr:
check("notification 202", rr.status == 202 and rr.read() == b"", "status=%s" % rr.status)
# 17. OPTIONS preflight -> 204 with CORS headers
req = urllib.request.Request("http://127.0.0.1:18692/mcp", method="OPTIONS")
try:
with urllib.request.urlopen(req, timeout=5) as rr:
check("options 204", rr.status == 204
and rr.headers.get("Access-Control-Allow-Origin") == "*",
"status=%s" % rr.status)
except urllib.error.HTTPError as e:
check("options 204", e.code == 204
and e.headers.get("Access-Control-Allow-Origin") == "*",
"status=%s" % e.code)
print("=== MCP ENDPOINT OK ===")
bridge_http_server.stop_server()
tests/test_sub_no_ws.py
+69
View File
@@ -0,0 +1,69 @@
# -*- coding: utf-8 -*-
"""Test: subscribe WITHOUT ws connection - does QMT push continue?"""
import sys, io, os, json, time, urllib.request
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src"))
if sys.stdout is None:
sys.stdout = io.StringIO()
import bridge_util, bridge_http_server, bridge_subscription
class FakeCtx:
def __init__(self):
self.subs = {}
self._next = 1
self.callback = None
self.push_calls = [] # record every QMT callback invocation
def subscribe_whole_quote(self, codes, callback=None):
h = self._next
self._next += 1
self.subs[h] = list(codes)
self.callback = callback
return h
def unsubscribe_quote(self, h):
self.subs.pop(h, None)
fake = FakeCtx()
bridge_util.CTX = fake
bridge_http_server.PORT = 18692
bridge_http_server.TOKEN = ""
bridge_http_server.start_server()
time.sleep(0.5)
# 1. subscribe whole (NO ws connection established)
req = urllib.request.Request(
"http://127.0.0.1:18692/data/subscribe",
data=json.dumps({"type": "whole", "codes": ["SH", "SZ"]}).encode(),
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=5) as r:
print("subscribe:", r.read().decode())
# 2. simulate QMT pushes (no ws connected)
print("QMT subs before pushes:", fake.subs)
print("WS clients count: (none connected)")
for i in range(3):
fake.callback({"600000.SH": {"lastPrice": 10.0 + i, "volume": 100}})
print("QMT push %d fired (callback invoked)" % (i + 1))
# 3. verify subscription still active (QMT sub not torn down)
print("QMT subs after pushes:", fake.subs)
print("sub active:", bool(fake.subs))
# 4. unsubscribe
req = urllib.request.Request(
"http://127.0.0.1:18692/data/unsubscribe",
data=json.dumps({"sub_id": 1}).encode(),
headers={"Content-Type": "application/json"},
)
try:
with urllib.request.urlopen(req, timeout=5) as r:
print("unsubscribe:", r.read().decode())
except Exception as e:
# find actual sub_id from state
import bridge_subscription as bs
print("unsub err:", e, "active subs:", bs.get_subscriptions())
print("QMT subs after unsubscribe:", fake.subs)
bridge_http_server.stop_server()
print("=== DONE ===")
tests/test_subscription.py
+118
View File
@@ -0,0 +1,118 @@
# -*- coding: utf-8 -*-
"""Local mock test for subscription endpoints (subscribe/unsubscribe/tick)."""
import sys, io, os, json, time, threading, urllib.request, urllib.error
if sys.stdout is None:
sys.stdout = io.StringIO()
if sys.stderr is None:
sys.stderr = io.StringIO()
RESULT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sub_test.txt")
SRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src")
sys.path.insert(0, SRC)
def log(msg):
with open(RESULT, "a") as f:
f.write(msg + "\n")
import bridge_util, bridge_http_server, bridge_ws_server, bridge_subscription
# ---- Fake QMT ContextInfo ----
class FakeCtx:
def __init__(self):
self.subs = {} # handle -> codes
self._next_handle = 1
self.pushed = [] # records of (handle, data) pushed
def subscribe_whole_quote(self, codes, callback=None):
handle = self._next_handle
self._next_handle += 1
self.subs[handle] = list(codes)
# remember callback so tests can trigger pushes
self.callback = callback
return handle
def unsubscribe_quote(self, handle):
self.subs.pop(handle, None)
def get_full_tick(self, codes):
return {c: {"lastPrice": 10.0, "volume": 100} for c in codes}
fake = FakeCtx()
bridge_util.CTX = fake
bridge_util.ACCOUNT = "TEST"
bridge_http_server.PORT = 18690
bridge_http_server.TOKEN = ""
bridge_http_server.start_server()
time.sleep(0.6)
BASE = "http://127.0.0.1:18690"
def post(path, obj):
data = json.dumps(obj).encode("utf-8")
req = urllib.request.Request(BASE + path, data=data, method="POST",
headers={"Content-Type": "application/json"})
try:
with urllib.request.urlopen(req, timeout=5) as r:
return r.status, json.loads(r.read().decode())
except urllib.error.HTTPError as e:
try:
return e.code, json.loads(e.read().decode())
except Exception:
return e.code, None
def get(path):
try:
with urllib.request.urlopen(BASE + path, timeout=5) as r:
return r.status, json.loads(r.read().decode())
except urllib.error.HTTPError as e:
try:
return e.code, json.loads(e.read().decode())
except Exception:
return e.code, None
# ---- Tests ----
# 1. subscribe whole
st, body = post("/data/subscribe", {"type": "whole", "codes": ["SH", "SZ"]})
log("1. subscribe whole -> HTTP %d body=%s" % (st, json.dumps(body, ensure_ascii=False)))
sub_id = body.get("sub_id") if body else None
assert st == 200 and sub_id is not None, "subscribe failed"
assert fake.subs, "QMT subscribe not created"
log(" QMT subs: %s" % fake.subs)
# 2. simulate QMT push -> should broadcast via ws (no ws client, no crash)
if fake.callback:
fake.callback({"600000.SH": {"lastPrice": 10.1}})
log("2. simulated QMT push OK (no crash)")
# 3. subscribe invalid type
st, body = post("/data/subscribe", {"type": "bogus", "codes": ["SH"]})
log("3. subscribe bogus type -> HTTP %d body=%s" % (st, json.dumps(body, ensure_ascii=False)))
assert st == 400, "bogus type should be 400"
# 4. subscribe missing codes
st, body = post("/data/subscribe", {"type": "whole"})
log("4. subscribe no codes -> HTTP %d body=%s" % (st, json.dumps(body, ensure_ascii=False)))
assert st == 400, "missing codes should be 400"
# 5. unsubscribe
st, body = post("/data/unsubscribe", {"sub_id": sub_id})
log("5. unsubscribe -> HTTP %d body=%s" % (st, json.dumps(body, ensure_ascii=False)))
assert st == 200, "unsubscribe failed"
assert not fake.subs, "QMT sub should be removed"
# 6. unsubscribe unknown
st, body = post("/data/unsubscribe", {"sub_id": 999999})
log("6. unsubscribe unknown -> HTTP %d body=%s" % (st, json.dumps(body, ensure_ascii=False)))
assert st == 404, "unknown sub should be 404"
# 7. GET /data/tick
st, body = get("/data/tick?codes=600000.SH,000001.SZ")
log("7. tick -> HTTP %d body=%s" % (st, json.dumps(body, ensure_ascii=False)))
assert st == 200 and "600000.SH" in body.get("data", {}), "tick failed"
# 8. tick missing codes
st, body = get("/data/tick")
log("8. tick no codes -> HTTP %d body=%s" % (st, json.dumps(body, ensure_ascii=False)))
assert st == 400, "tick no codes should be 400"
bridge_http_server.stop_server()
log("=== done ===")
tests/test_trade.py
+95
View File
@@ -0,0 +1,95 @@
# -*- coding: utf-8 -*-
"""Local test for trade endpoints (mock QMT). Run with QMT pythonw."""
import sys, io, os, json, time, urllib.request
if sys.stdout is None:
sys.stdout = io.StringIO()
if sys.stderr is None:
sys.stderr = io.StringIO()
RESULT = r"C:\Users\Docker\Development\qmt_bridge\tests\trade_test.txt"
SRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src")
sys.path.insert(0, SRC)
def log(msg):
with open(RESULT, "a") as f:
f.write(msg + "\n")
import bridge_util, bridge_http_server, bridge_data_adapter
class FakeCtx:
def get_market_data_ex(self, fields, codes, **kw):
import pandas as pd
df = pd.DataFrame({"open": [1.0], "close": [1.1]})
return {codes[0]: df}
def get_full_tick(self, codes):
return {c: {"lastPrice": 1.0} for c in codes}
def get_instrument_detail(self, code):
return {"InstrumentName": "贵州茅台", "UpStopPrice": 1438.67,
"DownStopPrice": 1177.09, "OpenDate": "20010827"}
def get_trading_dates(self, start="", end=""):
return ["20260818", "20260819", "20260820"]
bridge_util.CTX = FakeCtx()
bridge_util.ACCOUNT = "TESTACC"
def fake_trade(acct, atype, dtype):
if dtype == "position":
return [type("Pos", (), {"m_strInstrumentID": "600519", "m_strInstrumentName": "贵州茅台",
"m_nVolume": 100, "m_nCanUseVolume": 50,
"m_nFrozenVolume": 50, "m_dOpenPrice": 1500.0,
"m_dFloatProfit": 1000.0})()]
return [type("Pos", (), {"m_strInstrumentID": "600519.SH", "m_nVolume": 100,
"m_nCanUseVolume": 50, "m_dOpenPrice": 1500.0})()]
bridge_util.QMT_API["get_trade_detail_data"] = fake_trade
bridge_http_server.PORT = 18631
bridge_http_server.TOKEN = ""
bridge_http_server.ACCOUNT = "TESTACC"
# prime the trade cache by simulating a strategy-thread refresh
bridge_util._refresh_trade_cache("TESTACC", "STOCK")
bridge_http_server.start_server()
time.sleep(0.8)
BASE = "http://127.0.0.1:18631"
def get(path):
with urllib.request.urlopen(BASE + path, timeout=5) as r:
return json.loads(r.read().decode())
def check_positions():
"""Assert /trade/positions returns semantic fields + summary."""
r = get("/trade/positions")
data = r["data"]
assert r["ok"] is True, "positions ok flag"
assert "summary" in data, "summary present"
assert data["summary"]["count"] == 1, "summary count"
pos = data["positions"][0]
assert pos["stock_code"] == "600519.SH", "stock_code full"
assert pos["stock_name"] == "贵州茅台", "stock_name"
assert pos["volume"] == 100, "volume"
assert pos["available"] == 50, "available"
assert pos["avg_price"] == 1500.0, "avg_price"
assert pos["price"] == 1500.0, "price fallback to open"
assert pos["market_value"] == 150000.0, "market_value computed"
assert pos["profit"] == 1000.0, "profit from m_dFloatProfit"
assert pos["profit_pct"] == 0.67, "profit_pct computed"
# raw m_* fields preserved as superset
assert pos["m_strInstrumentID"] == "600519", "raw m_strInstrumentID"
assert pos["m_nVolume"] == 100, "raw m_nVolume"
log("OK positions enriched -> %s" % str(r)[:200])
for ep in ["/health", "/openapi.json", "/docs",
"/data/kline?code=600519.SH&period=1d&count=1",
"/data/quote?code=600519.SH", "/data/instrument?code=600519.SH",
"/data/calendar/trading_dates?start=20260818&end=20260820",
"/trade/positions", "/trade/asset",
"/trade/orders", "/trade/trades"]:
try:
if ep == "/trade/positions":
check_positions()
continue
r = get(ep)
log("OK %s -> %s" % (ep, str(r)[:80]))
except Exception as e:
log("FAIL %s -> %s" % (ep, str(e)[:80]))
bridge_http_server.stop_server()
log("=== done ===")
tests/test_ws_push.py
+86
View File
@@ -0,0 +1,86 @@
# -*- coding: utf-8 -*-
"""End-to-end test: WS connect -> subscribe -> QMT push -> WS receive."""
import sys, io, os, json, time, socket, struct, base64, urllib.request
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src"))
if sys.stdout is None:
sys.stdout = io.StringIO()
import bridge_util, bridge_http_server, bridge_subscription
class FakeCtx:
def __init__(self):
self.subs = {}
self._next = 1
self.callback = None
def subscribe_whole_quote(self, codes, callback=None):
h = self._next
self._next += 1
self.subs[h] = list(codes)
self.callback = callback
return h
def unsubscribe_quote(self, h):
self.subs.pop(h, None)
fake = FakeCtx()
bridge_util.CTX = fake
bridge_http_server.PORT = 18691
bridge_http_server.TOKEN = ""
bridge_http_server.start_server()
time.sleep(0.5)
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_recv(s, timeout=5):
s.settimeout(timeout)
hdr = s.recv(2)
if len(hdr) < 2:
return None
b1 = hdr[1]
ln = b1 & 0x7F
if ln == 126:
ln = struct.unpack(">H", s.recv(2))[0]
payload = b""
while len(payload) < ln:
chunk = s.recv(ln - len(payload))
if not chunk:
break
payload += chunk
return payload.decode("utf-8", "ignore")
# 1. connect WS
ws = ws_connect(18691)
print("ws connected")
# 2. subscribe whole via HTTP
req = urllib.request.Request(
"http://127.0.0.1:18691/data/subscribe",
data=json.dumps({"type": "whole", "codes": ["SH", "SZ"]}).encode(),
headers={"Content-Type": "application/json"},
)
with urllib.request.urlopen(req, timeout=5) as r:
print("subscribe:", r.read().decode())
# 3. trigger QMT push
time.sleep(0.3)
fake.callback({"600000.SH": {"lastPrice": 10.5, "volume": 200}})
# 4. WS should receive whole message
msg = ws_recv(ws)
print("ws received:", msg)
assert msg is not None, "ws received nothing!"
parsed = json.loads(msg)
assert parsed.get("type") == "whole", "expected type=whole, got %s" % parsed.get("type")
assert "600000.SH" in parsed.get("data", {}), "data missing code"
print("=== WS PUSH CHAIN OK ===")
ws.close()
bridge_http_server.stop_server()
tests/test_ws_real.py
+94
View File
@@ -0,0 +1,94 @@
# -*- 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)