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

import json

import pytest
import tornado.httpclient


class FakeOpenCodeClient:
    """Drop-in replacement for OpenCodeClient with recording + canned responses."""

    def __init__(self):
        self.calls = []
        self.session_id = "fake-session-123"
        self.health_response = {"healthy": True, "version": "0.1.0"}
        self.providers_response = {"providers": [{"id": "anthropic", "models": [{"id": "claude"}]}]}
        self.message_response = {
            "info": {"id": "msg-1"},
            "parts": [{"type": "text", "text": "def foo():\n    return 42\n"}],
        }
        # Canned session-messages list response (list_session_messages).
        # Default: one user + one assistant message, mixed text parts.
        self.messages_response = [
            {
                "info": {"role": "user", "id": "m1"},
                "parts": [{"type": "text", "text": "fix the bug"}],
            },
            {
                "info": {"role": "assistant", "id": "m2"},
                "parts": [{"type": "text", "text": "```python\nx = 1\n```"}],
            },
        ]

    @property
    def endpoint(self):
        return "http://fake-opencode"

    async def health(self):
        self.calls.append(("health",))
        return self.health_response

    async def list_providers(self):
        self.calls.append(("list_providers",))
        return self.providers_response

    async def create_session(self, title):
        self.calls.append(("create_session", title))
        return {"id": self.session_id, "title": title}

    async def send_message_sync(self, session_id, parts, provider_id=None, model_id=None, system=None):
        self.calls.append(("send_message_sync", session_id, parts, system))
        return self.message_response

    async def delete_session(self, session_id):
        self.calls.append(("delete_session", session_id))
        return True

    async def list_session_messages(self, session_id):
        self.calls.append(("list_session_messages", session_id))
        return self.messages_response


class FakeSessionManager:
    """Drop-in replacement for SessionManager with recording."""

    def __init__(self, session_id: str = "fake-session-123") -> None:
        self._session_id = session_id
        self.calls: list = []

    async def get_or_create(self, notebook_path: str) -> str:
        self.calls.append(("get_or_create", notebook_path))
        return self._session_id

    def peek(self, notebook_path: str):
        # Mirror SessionManager.peek: return the session id without
        # creating one. None means "no session yet" (used by the no-session
        # test to short-circuit the history route).
        return self._session_id or None

    async def release(self, notebook_path: str) -> bool:
        self.calls.append(("release", notebook_path))
        return True

    def list_sessions(self) -> list:
        return [
            {"notebookPath": path, "sessionId": sid}
            for path, sid in sorted(self._sessions.items())
        ]

    def invalidate(self, notebook_path: str) -> bool:
        self.calls.append(("invalidate", notebook_path))
        return True


async def test_hello(jp_fetch):
    # When
    response = await jp_fetch("opencode-bridge", "hello")

    # Then
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {
            "data": (
                "Hello, world!"
                " This is the '/opencode-bridge/hello' endpoint."
                " Try visiting me in your browser!"
            ),
        }


async def test_health_handler(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    response = await jp_fetch("opencode-bridge", "health")
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload["ok"] is True
    assert payload["version"] == "0.1.0"
    assert ("health",) in fake.calls


async def test_providers_handler(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    response = await jp_fetch("opencode-bridge", "providers")
    assert response.code == 200
    payload = json.loads(response.body)
    assert "providers" in payload
    assert payload["providers"][0]["id"] == "anthropic"
    assert ("list_providers",) in fake.calls


async def test_edit_handler(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    fake_sm = FakeSessionManager(session_id="fake-session-123")
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    body = json.dumps({
        "prompt": "Add type hints",
        "context": {
            "notebookPath": "test.ipynb",
            "cellId": "cell-1",
            "language": "python",
            "cellIndex": 0,
            "totalCells": 1,
            "source": "def foo(): return 42\n",
            "previousCode": None,
            "error": None,
        },
    })
    response = await jp_fetch(
        "opencode-bridge", "edit",
        method="POST",
        body=body,
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload["ok"] is True
    # Server now returns the AI reply as raw markdown (```fences``` kept
    # so the frontend marked.parse renders code blocks).
    assert payload["markdown"] == "def foo():\n    return 42"
    assert payload["sessionId"] == "fake-session-123"
    assert payload["notebookPath"] == "test.ipynb"

    # Session manager was used, not direct create/delete on client
    call_names = [c[0] for c in fake.calls]
    assert "create_session" not in call_names
    assert "delete_session" not in call_names
    assert "send_message_sync" in call_names

    # send_message_sync received the session ID from the manager
    send_call = [c for c in fake.calls if c[0] == "send_message_sync"][0]
    assert send_call[1] == "fake-session-123"
    # system prompt was passed (v3+ allows markdown/fences in the reply)
    assert "你是一个代码助手" in send_call[3]

    # SessionManager.get_or_create was called with the notebook path
    sm_call_names = [c[0] for c in fake_sm.calls]
    assert sm_call_names == ["get_or_create"]
    assert fake_sm.calls[0][1] == "test.ipynb"


async def test_session_list_handler(monkeypatch, jp_fetch):
    fake_sm = FakeSessionManager()
    fake_sm._sessions = {
        "foo.ipynb": "sid-1",
        "bar.ipynb": "sid-2",
    }  # direct injection
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    response = await jp_fetch("opencode-bridge", "sessions")
    assert response.code == 200
    payload = json.loads(response.body)
    assert "sessions" in payload
    paths = {s["notebookPath"] for s in payload["sessions"]}
    assert paths == {"foo.ipynb", "bar.ipynb"}


async def test_session_messages_handler_returns_empty_when_no_session(
    monkeypatch, jp_fetch
) -> None:
    # No session registered for this notebook -> handler returns
    # {"messages": []} WITHOUT calling OpenCodeClient (peek short-circuits).
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    fake_sm = FakeSessionManager(session_id="")
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    response = await jp_fetch(
        "opencode-bridge", "session-messages",
        method="GET",
        params={"notebook": "fresh.ipynb"},
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {"messages": []}
    # No OpenCode call was made.
    assert all(c[0] != "list_session_messages" for c in fake.calls)


async def test_session_messages_handler_projects_opencode_messages(
    monkeypatch, jp_fetch
) -> None:
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    fake_sm = FakeSessionManager(session_id="fake-session-123")
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    response = await jp_fetch(
        "opencode-bridge", "session-messages",
        method="GET",
        params={"notebook": "test.ipynb"},
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload == {
        "messages": [
            {"role": "user", "content": "fix the bug"},
            {"role": "assistant", "content": "```python\nx = 1\n```"},
        ]
    }
    # The OpenCode client was called with the session id from the manager.
    assert ("list_session_messages", "fake-session-123") in fake.calls


async def test_session_messages_handler_requires_notebook(
    monkeypatch, jp_fetch
) -> None:
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)
    # jp_fetch raises HTTPClientError on 4xx; assert the handler 400s
    # (the body would contain "missing 'notebook' query parameter").
    with pytest.raises(tornado.httpclient.HTTPClientError) as exc_info:
        await jp_fetch("opencode-bridge", "session-messages", method="GET")
    assert exc_info.value.code == 400


async def test_session_release_handler(monkeypatch, jp_fetch):
    fake = FakeOpenCodeClient()
    monkeypatch.setattr("opencode_bridge.routes.make_client", lambda h: fake)

    fake_sm = FakeSessionManager()
    monkeypatch.setattr(
        "opencode_bridge.routes.get_session_manager", lambda h: fake_sm
    )

    response = await jp_fetch(
        "opencode-bridge", "session",
        method="DELETE",
        params={"notebook": "foo.ipynb"},
    )
    assert response.code == 200
    payload = json.loads(response.body)
    assert payload["ok"] is True
    assert payload["notebookPath"] == "foo.ipynb"
    assert payload["deleted"] is True
    assert ("release", "foo.ipynb") in fake_sm.calls