iLSP/tests/test_ws_proxy.py

"""Integration tests: fake LSP client → WebSocket proxy → real pylsp.

Tests that a real LSP client (us, acting as the editor) can send LSP
JSON-RPC through the WebSocket proxy and get valid responses from the
actual pylsp process running inside the container.
"""

import asyncio
import json
import socket
import subprocess
import sys

import pytest
import aiohttp
from aiohttp import web

from ilsp.server import _build_app


# ---------------------------------------------------------------------------
# Helpers: LSP framing (Content-Length protocol)
# ---------------------------------------------------------------------------

def _lsp_frame(body: bytes) -> bytes:
    return b"Content-Length: " + str(len(body)).encode() + b"\r\n\r\n" + body


def _lsp_message(**kwargs) -> bytes:
    return _lsp_frame(json.dumps(kwargs).encode())


def _parse_lsp_frames(data: bytes) -> list[dict]:
    """Extract all complete LSP JSON-RPC messages from a raw byte buffer.

    pylsp sends multiple headers (Content-Length + Content-Type), so we
    parse each header line individually instead of assuming one header.
    """
    messages = []
    while b"\r\n\r\n" in data:
        headers_raw, rest = data.split(b"\r\n\r\n", 1)
        length = None
        for line in headers_raw.split(b"\r\n"):
            if line.lower().startswith(b"content-length:"):
                length = int(line.split(b":", 1)[1].strip())
                break
        if length is None or len(rest) < length:
            break
        messages.append(json.loads(rest[:length]))
        data = rest[length:]
    return messages


def _free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


# ---------------------------------------------------------------------------
# Fixtures: real pylsp + aiohttp proxy app
# ---------------------------------------------------------------------------

@pytest.fixture
async def pylsp_port():
    """Start real pylsp on a free TCP port; yield port; kill on teardown."""
    port = _free_port()
    proc = subprocess.Popen(
        [sys.executable, "-m", "pylsp", "--tcp", "--host", "127.0.0.1", "--port", str(port)],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    # Wait until pylsp is actually accepting connections
    for _ in range(20):
        await asyncio.sleep(0.1)
        try:
            r, w = await asyncio.open_connection("127.0.0.1", port)
            w.close()
            await w.wait_closed()
            break
        except OSError:
            continue
    yield port
    proc.terminate()
    proc.wait(timeout=5)


@pytest.fixture
async def proxy_base_url(monkeypatch, pylsp_port):
    """Start the aiohttp proxy pointing at real pylsp; yield base URL."""
    http_port = _free_port()
    monkeypatch.setattr("ilsp.server.PYTHON_LSP_PORT", pylsp_port)
    monkeypatch.setattr("ilsp.server.BICEP_LSP_PORT", _free_port())  # unused in these tests

    app = await _build_app()
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "127.0.0.1", http_port)
    await site.start()

    yield f"http://127.0.0.1:{http_port}"

    await runner.cleanup()


# ---------------------------------------------------------------------------
# Tests: fake LSP client talks to real pylsp via WebSocket
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_health_endpoint(proxy_base_url):
    """Sanity check: /health returns ok."""
    async with aiohttp.ClientSession() as session:
        async with session.get(f"{proxy_base_url}/health") as resp:
            assert resp.status == 200
            data = await resp.json()
            assert data["status"] == "ok"


@pytest.mark.asyncio
async def test_fake_client_initialize_gets_real_pylsp_capabilities(proxy_base_url):
    """
    Fake editor client sends LSP initialize through WebSocket proxy.
    Expects a real InitializeResult with pylsp capabilities back.
    """
    ws_url = proxy_base_url.replace("http://", "ws://") + "/python"
    init_request = _lsp_message(
        jsonrpc="2.0",
        id=1,
        method="initialize",
        params={
            "processId": None,
            "rootUri": None,
            "capabilities": {},
        },
    )

    async with aiohttp.ClientSession() as session:
        async with session.ws_connect(ws_url) as ws:
            await ws.send_bytes(init_request)

            # Collect response bytes — pylsp may send in multiple frames
            buf = b""
            for _ in range(5):
                msg = await asyncio.wait_for(ws.receive(), timeout=5.0)
                if msg.type not in (aiohttp.WSMsgType.BINARY, aiohttp.WSMsgType.TEXT):
                    break
                buf += msg.data if msg.type == aiohttp.WSMsgType.BINARY else msg.data.encode()
                if b"\r\n\r\n" in buf:
                    break

            messages = _parse_lsp_frames(buf)
            assert messages, "No LSP response received from pylsp"

            # Find the InitializeResult (id=1)
            response = next((m for m in messages if m.get("id") == 1), None)
            assert response is not None, f"No response with id=1, got: {messages}"
            assert "result" in response, f"Expected result, got: {response}"
            assert "capabilities" in response["result"], "pylsp did not return capabilities"


@pytest.mark.asyncio
async def test_fake_client_receives_completion_capabilities(proxy_base_url):
    """pylsp should advertise completionProvider in its capabilities."""
    ws_url = proxy_base_url.replace("http://", "ws://") + "/python"
    init_request = _lsp_message(
        jsonrpc="2.0",
        id=1,
        method="initialize",
        params={"processId": None, "rootUri": None, "capabilities": {}},
    )

    async with aiohttp.ClientSession() as session:
        async with session.ws_connect(ws_url) as ws:
            await ws.send_bytes(init_request)

            buf = b""
            for _ in range(5):
                msg = await asyncio.wait_for(ws.receive(), timeout=5.0)
                if msg.type not in (aiohttp.WSMsgType.BINARY, aiohttp.WSMsgType.TEXT):
                    break
                buf += msg.data if msg.type == aiohttp.WSMsgType.BINARY else msg.data.encode()
                if b"\r\n\r\n" in buf:
                    break

            messages = _parse_lsp_frames(buf)
            response = next((m for m in messages if m.get("id") == 1), None)
            caps = response["result"]["capabilities"]
            assert "completionProvider" in caps, f"No completionProvider in: {caps}"


@pytest.mark.asyncio
async def test_ws_backend_unavailable_closes_cleanly(proxy_base_url, monkeypatch):
    """If pylsp is unreachable, WebSocket should close gracefully (not crash)."""
    # Use _free_port() to guarantee nothing is listening (avoids CI port 19999 collisions)
    dead_port = _free_port()
    monkeypatch.setattr("ilsp.server.PYTHON_LSP_PORT", dead_port)

    ws_url = proxy_base_url.replace("http://", "ws://") + "/python"
    async with aiohttp.ClientSession() as session:
        async with session.ws_connect(ws_url) as ws:
            msg = await asyncio.wait_for(ws.receive(), timeout=5.0)
            assert msg.type in (
                aiohttp.WSMsgType.CLOSE,
                aiohttp.WSMsgType.CLOSED,
                aiohttp.WSMsgType.ERROR,
            )