Docker
87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
# -*- 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()
|