notebook-ai-extension/opencode_bridge/tests/test_client.py

"""Tests for opencode_bridge.opencode_client."""

import json
from io import BytesIO
from unittest.mock import MagicMock

import pytest
import tornado.httpclient
import tornado.httputil

from opencode_bridge.config import OpenCodeConfig
from opencode_bridge.opencode_client import OpenCodeClient, OpenCodeError


class MockHTTPClient:
    def __init__(self) -> None:
        self.calls: list[dict] = []
        self.responses: list[tuple[int, object]] = []

    async def fetch(self, request, **kwargs) -> tornado.httpclient.HTTPResponse:
        self.calls.append({"request": request, "kwargs": kwargs})
        status, body = self.responses.pop(0)
        buffer = BytesIO(json.dumps(body).encode("utf-8"))
        return tornado.httpclient.HTTPResponse(
            request=request,
            code=status,
            headers=tornado.httputil.HTTPHeaders(),
            buffer=buffer,
        )


@pytest.fixture
def base_config() -> OpenCodeConfig:
    return OpenCodeConfig(
        url="http://127.0.0.1:4096",
        user="opencode",
        password="",
        request_timeout_seconds=120,
    )


def _make_client(
    config: OpenCodeConfig, responses: list[tuple[int, object]]
) -> tuple[OpenCodeClient, MockHTTPClient]:
    mock = MockHTTPClient()
    mock.responses = responses
    return OpenCodeClient(config, http_client=mock), mock


@pytest.mark.asyncio
async def test_health_gets_global_health(base_config: OpenCodeConfig) -> None:
    client, mock = _make_client(base_config, [(200, {"status": "ok"})])
    result = await client.health()
    assert result == {"status": "ok"}
    assert len(mock.calls) == 1
    call = mock.calls[0]
    assert call["request"].url == "http://127.0.0.1:4096/global/health"
    assert call["request"].method == "GET"


@pytest.mark.asyncio
async def test_create_session_posts_title(base_config: OpenCodeConfig) -> None:
    client, mock = _make_client(base_config, [(200, {"id": "s1"})])
    result = await client.create_session("foo")
    assert result == {"id": "s1"}
    call = mock.calls[0]
    assert call["request"].url == "http://127.0.0.1:4096/session"
    assert call["request"].method == "POST"
    body = json.loads(call["request"].body.decode("utf-8"))
    assert body == {"title": "foo"}


@pytest.mark.asyncio
async def test_auth_set_on_request_when_password_present(
    base_config: OpenCodeConfig,
) -> None:
    """Bug fix: auth must live on the HTTPRequest, not as fetch kwargs
    (tornado rejects fetch(request, **kwargs) for request-construction kwargs)."""
    config = base_config._replace(password="secret")
    client, mock = _make_client(config, [(200, {"status": "ok"})])
    await client.health()
    call = mock.calls[0]
    # No kwargs forwarded to fetch.
    assert call["kwargs"] == {}
    # Auth is on the HTTPRequest.
    assert call["request"].auth_username == "opencode"
    assert call["request"].auth_password == "secret"


@pytest.mark.asyncio
async def test_auth_absent_when_no_password(base_config: OpenCodeConfig) -> None:
    client, mock = _make_client(base_config, [(200, {"status": "ok"})])
    await client.health()
    call = mock.calls[0]
    assert call["kwargs"] == {}
    assert call["request"].auth_username is None
    assert call["request"].auth_password is None


@pytest.mark.asyncio
async def test_send_message_async_includes_model_when_provider_and_model_given(
    base_config: OpenCodeConfig,
) -> None:
    """send_message_async posts to /prompt_async with provider/model when both set."""
    client, mock = _make_client(base_config, [(200, {"done": True})])
    parts = [{"type": "text", "text": "hello"}]
    result = await client.send_message_async("s1", parts, provider_id="p1", model_id="m1")
    assert result is None
    call = mock.calls[0]
    assert call["request"].url == "http://127.0.0.1:4096/session/s1/prompt_async"
    assert call["request"].method == "POST"
    body = json.loads(call["request"].body.decode("utf-8"))
    assert body["parts"] == parts
    assert body["model"] == {"providerID": "p1", "modelID": "m1"}


@pytest.mark.asyncio
async def test_send_message_async_omits_model_when_provider_or_model_missing(
    base_config: OpenCodeConfig,
) -> None:
    client, mock = _make_client(base_config, [(200, {"done": True})])
    parts = [{"type": "text", "text": "hello"}]
    result = await client.send_message_async("s1", parts)
    assert result is None
    call = mock.calls[0]
    body = json.loads(call["request"].body.decode("utf-8"))
    assert body["parts"] == parts
    assert "model" not in body


@pytest.mark.asyncio
async def test_send_message_async_includes_system_when_provided() -> None:
    """system param is forwarded into the request body when not None."""
    config = OpenCodeConfig(
        url="http://x:1", user="u", password="", request_timeout_seconds=120
    )
    mock = MockHTTPClient()
    mock.responses.append((200, {"info": {}, "parts": []}))
    client = OpenCodeClient(config, http_client=mock)

    await client.send_message_async("sid", [{"type": "text", "text": "hi"}], system="be brief")

    assert len(mock.calls) == 1
    body = json.loads(mock.calls[0]["request"].body.decode("utf-8"))
    assert body["system"] == "be brief"
    assert body["parts"] == [{"type": "text", "text": "hi"}]


@pytest.mark.asyncio
async def test_stream_global_events_yields_parsed_events(
    base_config: OpenCodeConfig,
) -> None:
    """stream_global_events yields one dict per `data: {...}` SSE block."""

    class _StreamingMock(MockHTTPClient):
        async def fetch(self, request, **kwargs):
            self.calls.append({"request": request, "kwargs": kwargs})
            # Simulate OpenCode's /global/event: three events back-to-back.
            chunks = [
                b"data: {\"type\": \"session.next.text.delta\", \"properties\": "
                b"{\"sessionID\": \"s1\", \"delta\": \"hello \"}}\n\n",
                b"data: {\"type\": \"session.next.text.delta\", \"properties\": "
                b"{\"sessionID\": \"s1\", \"delta\": \"world\"}}\n\n",
                b"data: {\"type\": \"session.idle\", \"properties\": "
                b"{\"sessionID\": \"s1\"}}\n\n",
            ]
            # streaming_callback now lives on the HTTPRequest (per tornado
            # API — fetch() rejects kwargs when an HTTPRequest is passed).
            cb = getattr(request, "streaming_callback", None)
            if cb is not None:
                for c in chunks:
                    cb(c)
            # Empty body, code 200 to terminate normally.
            return tornado.httpclient.HTTPResponse(
                request=request,
                code=200,
                headers=tornado.httputil.HTTPHeaders(),
                buffer=BytesIO(b""),
            )

    mock = _StreamingMock()
    client = OpenCodeClient(base_config, http_client=mock)

    events = []
    async for ev in client.stream_global_events():
        events.append(ev)

    assert len(events) == 3
    assert events[0]["type"] == "session.next.text.delta"
    assert events[0]["properties"]["delta"] == "hello "
    assert events[1]["properties"]["delta"] == "world"
    assert events[2]["type"] == "session.idle"
    # Stream used the SSE accept header and no timeout.
    call = mock.calls[0]
    assert call["kwargs"] == {"raise_error": False}
    assert call["request"].method == "GET"
    assert call["request"].url == "http://127.0.0.1:4096/global/event"
    assert call["request"].request_timeout == 0
    assert call["request"].headers.get("Accept") == "text/event-stream"
    assert call["request"].streaming_callback is not None


def test_parse_sse_block_handles_concatenated_data_lines() -> None:
    """Multi-line `data:` blocks are joined with \n (per SSE spec)."""
    parsed = OpenCodeClient._parse_sse_block(
        "data: {\"a\": 1,\ndata:  \"b\": 2}"
    )
    assert parsed == {"a": 1, "b": 2}


def test_parse_sse_block_ignores_non_data_fields() -> None:
    parsed = OpenCodeClient._parse_sse_block("event: ping\ndata: {\"k\": \"v\"}\n\n")
    assert parsed == {"k": "v"}


def test_parse_sse_block_returns_none_on_empty() -> None:
    assert OpenCodeClient._parse_sse_block("") is None
    assert OpenCodeClient._parse_sse_block("event: ping\n\n") is None
    assert OpenCodeClient._parse_sse_block(": heartbeat\n\n") is None


def test_parse_sse_block_returns_none_on_bad_json() -> None:
    assert OpenCodeClient._parse_sse_block("data: not-json\n\n") is None


@pytest.mark.asyncio
async def test_delete_session_returns_bool_by_status(base_config: OpenCodeConfig) -> None:
    client, mock = _make_client(base_config, [(200, True), (404, False)])
    assert await client.delete_session("s1") is True
    assert await client.delete_session("s1") is False
    assert len(mock.calls) == 2
    assert mock.calls[0]["request"].method == "DELETE"
    assert mock.calls[0]["request"].url == "http://127.0.0.1:4096/session/s1"


@pytest.mark.asyncio
async def test_reply_permission_posts_response(base_config: OpenCodeConfig) -> None:
    """reply_permission forwards the response string to /session/:id/permissions/:permId."""
    client, mock = _make_client(base_config, [(200, True)])
    ok = await client.reply_permission("s1", "perm-1", "once")
    assert ok is True
    call = mock.calls[0]
    assert call["request"].url == "http://127.0.0.1:4096/session/s1/permissions/perm-1"
    assert call["request"].method == "POST"
    body = json.loads(call["request"].body.decode("utf-8"))
    assert body == {"response": "once"}


@pytest.mark.asyncio
async def test_reply_permission_returns_false_on_404(
    base_config: OpenCodeConfig,
) -> None:
    client, mock = _make_client(base_config, [(404, False)])
    assert await client.reply_permission("s1", "perm-1", "reject") is False


@pytest.mark.asyncio
async def test_reply_question_posts_answer(base_config: OpenCodeConfig) -> None:
    """reply_question forwards the answer to /session/:id/question/:qid/reply."""
    client, mock = _make_client(base_config, [(200, True)])
    ok = await client.reply_question("s1", "q-7", "use 3.12")
    assert ok is True
    call = mock.calls[0]
    assert call["request"].url == "http://127.0.0.1:4096/session/s1/question/q-7/reply"
    body = json.loads(call["request"].body.decode("utf-8"))
    assert body == {"answer": "use 3.12"}


@pytest.mark.asyncio
async def test_list_all_sessions_gets_session(base_config: OpenCodeConfig) -> None:
    """list_all_sessions proxies OpenCode Serve's GET /session and returns
    the raw list (no projection)."""
    payload = [
        {"id": "s1", "title": "first", "createdAt": 1, "updatedAt": 2},
        {"id": "s2", "title": "second", "createdAt": 3, "updatedAt": 4},
    ]
    client, mock = _make_client(base_config, [(200, payload)])
    sessions = await client.list_all_sessions()
    assert sessions == payload
    call = mock.calls[0]
    assert call["request"].url == "http://127.0.0.1:4096/session"
    assert call["request"].method == "GET"


@pytest.mark.asyncio
async def test_non_2xx_response_raises_with_body(base_config: OpenCodeConfig) -> None:
    client, mock = _make_client(base_config, [(500, {"error": "boom"})])
    with pytest.raises(OpenCodeError) as exc_info:
        await client.health()
    assert "boom" in str(exc_info.value)


def test_close_is_idempotent() -> None:
    """close() can be called multiple times safely."""
    client = OpenCodeClient(OpenCodeConfig(
        url="http://x", user="u", password="", request_timeout_seconds=10
    ))
    client.close()
    client.close()  # must not raise
    assert client._http_client is None


def test_close_does_not_close_injected_client() -> None:
    """close() must not call .close() on a client it didn't create."""
    injected = MagicMock()
    client = OpenCodeClient(OpenCodeConfig(
        url="http://x", user="u", password="", request_timeout_seconds=10
    ), http_client=injected)
    client.close()
    injected.close.assert_not_called()
    assert client._http_client is None