tsl-devkit/lsp-server/test/test_core_server.py

#!/usr/bin/env python3

import argparse
import json
import os
import select
import subprocess
import time
from pathlib import Path


def frame(message: dict) -> bytes:
    body = json.dumps(message, separators=(",", ":")).encode("utf-8")
    return f"Content-Length: {len(body)}\r\n\r\n".encode("ascii") + body


def read_messages(data: bytes) -> list[dict]:
    messages = []
    offset = 0
    while offset < len(data):
        header_end = data.find(b"\r\n\r\n", offset)
        if header_end < 0:
            raise RuntimeError("Incomplete LSP response header")

        header = data[offset:header_end].decode("ascii")
        fields = dict(line.split(": ", 1) for line in header.split("\r\n"))
        length = int(fields["Content-Length"])
        body_start = header_end + 4
        body_end = body_start + length
        if body_end > len(data):
            raise RuntimeError("Incomplete LSP response body")

        messages.append(json.loads(data[body_start:body_end].decode("utf-8")))
        offset = body_end
    return messages


def response_by_id(data: bytes, request_id: int | str) -> dict:
    return next(message for message in read_messages(data)
                if message.get("id") == request_id)


def run_raw(server: Path, payload: bytes) -> subprocess.CompletedProcess[bytes]:
    return subprocess.run(
        [str(server), "--core-server-fixture"],
        input=payload,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=5,
        check=False,
    )


def run_batch(server: Path, messages: list[dict]) -> subprocess.CompletedProcess[bytes]:
    return run_raw(server, b"".join(frame(message) for message in messages))


class LspClient:
    def __init__(self, server: Path) -> None:
        self.proc = subprocess.Popen(
            [str(server), "--core-server-fixture"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        self.buffer = b""

    def send(self, message: dict) -> None:
        if self.proc.stdin is None:
            raise RuntimeError("Fixture stdin is unavailable")
        self.proc.stdin.write(frame(message))
        self.proc.stdin.flush()

    def read(self, timeout: float = 1.0) -> dict:
        if self.proc.stdout is None:
            raise RuntimeError("Fixture stdout is unavailable")

        deadline = time.monotonic() + timeout
        while True:
            header_end = self.buffer.find(b"\r\n\r\n")
            if header_end >= 0:
                header = self.buffer[:header_end].decode("ascii")
                fields = dict(line.split(": ", 1)
                              for line in header.split("\r\n"))
                length = int(fields["Content-Length"])
                body_start = header_end + 4
                body_end = body_start + length
                if len(self.buffer) >= body_end:
                    body = self.buffer[body_start:body_end]
                    self.buffer = self.buffer[body_end:]
                    return json.loads(body.decode("utf-8"))

            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError("Timed out waiting for LSP response")

            ready, _, _ = select.select(
                [self.proc.stdout.fileno()], [], [], remaining)
            if not ready:
                continue

            chunk = os.read(self.proc.stdout.fileno(), 4096)
            if not chunk:
                raise RuntimeError("Fixture closed stdout before a response")
            self.buffer += chunk

    def close_input(self) -> int:
        if self.proc.stdin is not None:
            self.proc.stdin.close()
        return self.proc.wait(timeout=5)

    def kill(self) -> None:
        if self.proc.poll() is None:
            self.proc.kill()
            self.proc.wait(timeout=5)


def assert_lifecycle(server: Path) -> None:
    before_init = run_batch(server, [
        {"jsonrpc": "2.0", "id": 1, "method": "shutdown"},
        {"jsonrpc": "2.0", "method": "exit"},
    ])
    if before_init.returncode != 1:
        raise RuntimeError("shutdown before initialize should exit with code 1")
    if response_by_id(before_init.stdout, 1)["error"]["code"] != -32002:
        raise RuntimeError("shutdown before initialize should return ServerNotInitialized")

    repeated = run_batch(server, [
        {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}},
        {"jsonrpc": "2.0", "id": 2, "method": "initialize", "params": {}},
        {"jsonrpc": "2.0", "id": 3, "method": "shutdown"},
        {"jsonrpc": "2.0", "method": "exit"},
    ])
    if repeated.returncode != 0:
        raise RuntimeError("shutdown followed by exit should return code 0")
    if response_by_id(repeated.stdout, 2)["error"]["code"] != -32600:
        raise RuntimeError("repeated initialize should return InvalidRequest")
    if response_by_id(repeated.stdout, 3).get("result", "missing") is not None:
        raise RuntimeError("shutdown should return a null result")

    direct_exit = run_batch(
        server, [{"jsonrpc": "2.0", "method": "exit"}])
    if direct_exit.returncode != 1:
        raise RuntimeError("exit before shutdown should return code 1")

    client = LspClient(server)
    try:
        client.send({
            "jsonrpc": "2.0",
            "id": 11,
            "method": "initialize",
            "params": {},
        })
        if client.read().get("id") != 11:
            raise RuntimeError("missing initialize response")

        client.send({"jsonrpc": "2.0", "id": 12, "method": "shutdown"})
        if client.read().get("result", "missing") is not None:
            raise RuntimeError("shutdown should return a null result")

        time.sleep(0.1)
        if client.proc.poll() is not None:
            raise RuntimeError("server exited before receiving exit notification")

        client.send({"jsonrpc": "2.0", "method": "exit"})
        if client.close_input() != 0:
            raise RuntimeError("orderly lifecycle should return code 0")
    finally:
        client.kill()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", type=Path, required=True)
    args = parser.parse_args()

    assert_lifecycle(args.server)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())