Test

Author: krassowski

test/test_python_lsp.py

@@ -0,0 +1,118 @@
+import socket
+import threading
+import base64
+import hashlib
+import json
+import struct
+import pytest
+import time
+from pylsp.python_lsp import start_ws_lang_server, PythonLSPServer
+
+WS_PORT = 3002
+NUM_REQUESTS = 20  # enough to provoke interleaving
+
+# --- Helpers for raw WebSocket framing ---
+
+def make_ws_key():
+    raw = hashlib.sha1(str(time.time()).encode()).digest()[:16]
+    return base64.b64encode(raw).decode()
+
+
+def build_handshake(host='localhost', port=WS_PORT):
+    key = make_ws_key()
+    return (
+        f"GET / HTTP/1.1\r\n"
+        f"Host: {host}:{port}\r\n"
+        f"Upgrade: websocket\r\n"
+        f"Connection: Upgrade\r\n"
+        f"Sec-WebSocket-Key: {key}\r\n"
+        f"Sec-WebSocket-Version: 13\r\n"
+        f"\r\n"
+    ).encode()
+
+
+def build_frame(text: str) -> bytes:
+    data = text.encode('utf8')
+    length = len(data)
+    header = b"\x81"
+    if length < 126:
+        header += struct.pack('B', length)
+    elif length < (1 << 16):
+        header += b"\x7e" + struct.pack('!H', length)
+    else:
+        header += b"\x7f" + struct.pack('!Q', length)
+    return header + data
+
+
+# Helper to run the LSP server in a background thread
+def run_server():
+    # Use the actual PythonLSPServer class for handler_class
+    start_ws_lang_server(host='127.0.0.1', port=WS_PORT, check_parent_process=False, handler_class=PythonLSPServer)
+
+@pytest.fixture(scope="module", autouse=True)
+def server():
+    # Start the server thread
+    thread = threading.Thread(target=run_server, daemon=True)
+    thread.start()
+    # Give server time to start
+    time.sleep(5)
+    yield
+    # Daemon thread will exit with test process
+
+
+def test_raw_frame_json_integrity():
+    # 1) Open raw TCP socket and perform WebSocket handshake
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.connect(('127.0.0.1', WS_PORT))
+    sock.sendall(build_handshake())
+    resp = b''
+    while b'\r\n\r\n' not in resp:
+        resp += sock.recv(1024)
+
+    # 2) Send frames concurrently
+    def send_req(i):
+        msg = {"jsonrpc":"2.0","id":i,"method":"initialize","params":{}}
+        sock.sendall(build_frame(json.dumps(msg, ensure_ascii=False)))
+
+    threads = []
+    for i in range(1, NUM_REQUESTS + 1):
+        t = threading.Thread(target=send_req, args=(i,))
+        t.start()
+        threads.append(t)
+    for t in threads:
+        t.join()
+
+    # 3) Read and validate responses
+    received_ids = set()
+    buffer = b''
+    while len(received_ids) < NUM_REQUESTS:
+        chunk = sock.recv(4096)
+        if not chunk:
+            break
+        buffer += chunk
+        # parse frames
+        while True:
+            if len(buffer) < 2:
+                break
+            b1, b2 = buffer[0], buffer[1]
+            length = b2 & 0x7f
+            offset = 2
+            if length == 126:
+                if len(buffer) < offset + 2: break
+                length = struct.unpack('!H', buffer[offset:offset+2])[0]; offset += 2
+            elif length == 127:
+                if len(buffer) < offset + 8: break
+                length = struct.unpack('!Q', buffer[offset:offset+8])[0]; offset += 8
+            if b2 & 0x80:
+                offset += 4
+            if len(buffer) < offset + length:
+                break
+            payload = buffer[offset:offset+length]
+            buffer = buffer[offset+length:]
+            try:
+                obj = json.loads(payload.decode('utf8'))
+            except json.JSONDecodeError:
+                pytest.fail(f"Invalid JSON frame: {payload}")
+            received_ids.add(obj.get('id'))
+    assert received_ids == set(range(1, NUM_REQUESTS + 1)), f"Expected IDs 1..{NUM_REQUESTS}, got {received_ids}"
+    sock.close()