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

"""Unit tests for SessionManager — no Jupyter server fixture."""
from __future__ import annotations

import pytest

from opencode_bridge.session_manager import SessionManager


class FakeClient:
    """Mimics OpenCodeClient just enough for SessionManager.

    By default returns the same id for every create_session call
    (matching the original legacy behavior; several pre-existing tests
    rely on that). Tests that need a fresh id per call should set
    ``unique_ids_per_call=True`` in the constructor.
    """
    def __init__(
        self,
        session_id: str = "sid-from-fake",
        unique_ids_per_call: bool = False,
    ) -> None:
        self._session_id = session_id
        self._unique_ids = unique_ids_per_call
        self._next_id = 0
        self.calls: list = []

    async def create_session(self, title: str) -> dict:
        self.calls.append(("create_session", title))
        if self._unique_ids:
            self._next_id += 1
            return {
                "id": "%s-%d" % (self._session_id, self._next_id),
                "title": title,
            }
        return {"id": self._session_id, "title": title}

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


@pytest.mark.asyncio
async def test_get_or_create_first_call_creates_session():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    sid = await sm.get_or_create("foo.ipynb")
    assert sid == "sid-from-fake"
    assert sm.has_session("foo.ipynb")
    assert client.calls == [("create_session", "jupyter:foo.ipynb")]


@pytest.mark.asyncio
async def test_get_or_create_second_call_returns_same_sid():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    sid1 = await sm.get_or_create("foo.ipynb")
    sid2 = await sm.get_or_create("foo.ipynb")
    assert sid1 == sid2
    # Only ONE create_session call total
    create_calls = [c for c in client.calls if c[0] == "create_session"]
    assert len(create_calls) == 1


@pytest.mark.asyncio
async def test_different_notebooks_get_different_sessions():
    client1 = FakeClient(session_id="sid-1")
    client2 = FakeClient(session_id="sid-2")
    factories = [client1, client2]
    sm = SessionManager(lambda: factories.pop(0) if factories else client1)
    sid1 = await sm.get_or_create("foo.ipynb")
    sid2 = await sm.get_or_create("bar.ipynb")
    assert sid1 == "sid-1"
    assert sid2 == "sid-2"


@pytest.mark.asyncio
async def test_release_returns_true_and_clears_session():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    await sm.get_or_create("foo.ipynb")
    deleted = await sm.release("foo.ipynb")
    assert deleted is True
    assert not sm.has_session("foo.ipynb")
    assert ("delete_session", "sid-from-fake") in client.calls


@pytest.mark.asyncio
async def test_release_returns_false_when_no_session():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    deleted = await sm.release("never-existed.ipynb")
    assert deleted is False
    assert client.calls == []


@pytest.mark.asyncio
async def test_get_or_create_after_release_creates_new_session():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    sid1 = await sm.get_or_create("foo.ipynb")
    await sm.release("foo.ipynb")
    sid2 = await sm.get_or_create("foo.ipynb")
    # Same fake client returns same ID, but two create_session calls happened
    assert sid1 == sid2
    create_calls = [c for c in client.calls if c[0] == "create_session"]
    assert len(create_calls) == 2


@pytest.mark.asyncio
async def test_invalidate_drops_cached_sid_without_calling_opencode():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    await sm.get_or_create("foo.ipynb")
    calls_before = len(client.calls)
    removed = sm.invalidate("foo.ipynb")
    assert removed is True
    assert not sm.has_session("foo.ipynb")
    # No additional calls to client
    assert len(client.calls) == calls_before


def test_list_sessions_empty():
    sm = SessionManager(lambda: FakeClient())
    assert sm.list_sessions() == []


@pytest.mark.asyncio
async def test_list_sessions_returns_all_mappings():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    await sm.get_or_create("a.ipynb")
    # Manually inject a second to test listing
    sm._active["b.ipynb"] = "sid-b"
    listing = sm.list_sessions()
    paths = {s["notebookPath"] for s in listing}
    assert paths == {"a.ipynb", "b.ipynb"}


@pytest.mark.asyncio
async def test_concurrent_get_or_create_does_not_double_create():
    """Two concurrent calls for the same path must share one session."""
    import asyncio

    client = FakeClient()
    sm = SessionManager(lambda: client)
    sids = await asyncio.gather(
        sm.get_or_create("foo.ipynb"),
        sm.get_or_create("foo.ipynb"),
        sm.get_or_create("foo.ipynb"),
    )
    assert sids[0] == sids[1] == sids[2]
    create_calls = [c for c in client.calls if c[0] == "create_session"]
    assert len(create_calls) == 1, "concurrent get_or_create must not double-create"


# ---------------------------------------------------------------------------
# Multi-session: 1 notebook can bind N sessions, one is active
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_create_new_binds_and_makes_active():
    client = FakeClient(unique_ids_per_call=True)
    sm = SessionManager(lambda: client)
    sid = await sm.create_new("foo.ipynb", title="custom")
    assert sid == "sid-from-fake-1"
    # New session is the active one and the only bound one.
    assert sm.get_active("foo.ipynb") == sid
    assert sm.list_sessions_for_notebook("foo.ipynb") == [sid]


@pytest.mark.asyncio
async def test_create_new_when_active_exists_keeps_old_bound():
    client = FakeClient(unique_ids_per_call=True)
    sm = SessionManager(lambda: client)
    first = await sm.get_or_create("foo.ipynb")
    second = await sm.create_new("foo.ipynb")
    # Two distinct session ids.
    assert first != second
    assert first == "sid-from-fake-1"
    assert second == "sid-from-fake-2"
    # Both are bound; the new one is active.
    bound = sm.list_sessions_for_notebook("foo.ipynb")
    assert set(bound) == {first, second}
    assert sm.get_active("foo.ipynb") == second
    # get_or_create still returns the active one (does NOT auto-pick the latest).
    assert await sm.get_or_create("foo.ipynb") == second


def test_bind_existing_does_not_call_opencode():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    sm.bind_existing("foo.ipynb", "sid-from-list")
    assert client.calls == []  # no OpenCode round-trip
    assert sm.get_active("foo.ipynb") == "sid-from-list"
    assert "sid-from-list" in sm.list_sessions_for_notebook("foo.ipynb")


def test_set_active_requires_session_to_be_bound():
    sm = SessionManager(lambda: FakeClient())
    sm.bind_existing("foo.ipynb", "sid-A")
    assert sm.set_active("foo.ipynb", "sid-A") is True
    # Switching to a session that is NOT bound must fail.
    assert sm.set_active("foo.ipynb", "sid-UNKNOWN") is False
    # Active didn't change.
    assert sm.get_active("foo.ipynb") == "sid-A"


def test_unbind_removes_from_bound_and_clears_active_if_was_active():
    sm = SessionManager(lambda: FakeClient())
    sm.bind_existing("foo.ipynb", "sid-A")
    sm.bind_existing("foo.ipynb", "sid-B")
    # B is the active one (last bind wins). Unbind B.
    assert sm.get_active("foo.ipynb") == "sid-B"
    assert sm.unbind("foo.ipynb", "sid-B") is True
    # Active falls back to None (no other session auto-promotes).
    assert sm.get_active("foo.ipynb") is None
    assert sm.list_sessions_for_notebook("foo.ipynb") == ["sid-A"]


def test_unbind_non_active_leaves_active_intact():
    sm = SessionManager(lambda: FakeClient())
    sm.bind_existing("foo.ipynb", "sid-A")
    sm.bind_existing("foo.ipynb", "sid-B")
    assert sm.get_active("foo.ipynb") == "sid-B"
    # Unbind A (not the active one).
    assert sm.unbind("foo.ipynb", "sid-A") is True
    # Active is still B.
    assert sm.get_active("foo.ipynb") == "sid-B"


def test_unbind_active_keeps_other_sessions():
    sm = SessionManager(lambda: FakeClient())
    sm.bind_existing("foo.ipynb", "sid-A")
    sm.bind_existing("foo.ipynb", "sid-B")
    sm.unbind_active("foo.ipynb")
    assert sm.get_active("foo.ipynb") is None
    # B is gone from the bound list; A remains.
    assert "sid-A" in sm.list_sessions_for_notebook("foo.ipynb")
    assert "sid-B" not in sm.list_sessions_for_notebook("foo.ipynb")


@pytest.mark.asyncio
async def test_delete_session_unbinds_and_calls_opencode():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    sm.bind_existing("foo.ipynb", "sid-A")
    sm.bind_existing("foo.ipynb", "sid-B")
    deleted = await sm.delete_session("foo.ipynb", "sid-A")
    assert deleted is True
    # OpenCode was called.
    assert any(
        c[0] == "delete_session" and c[1] == "sid-A" for c in client.calls
    )
    # sid-A is gone from bindings; sid-B remains as active.
    assert "sid-A" not in sm.list_sessions_for_notebook("foo.ipynb")
    assert sm.get_active("foo.ipynb") == "sid-B"


@pytest.mark.asyncio
async def test_delete_session_not_bound_returns_false():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    assert await sm.delete_session("foo.ipynb", "never-bound") is False
    # No OpenCode call.
    assert not any(c[0] == "delete_session" for c in client.calls)


@pytest.mark.asyncio
async def test_release_drops_all_bindings_and_deletes_active():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    sm.bind_existing("foo.ipynb", "sid-A")
    sm.bind_existing("foo.ipynb", "sid-B")
    ok = await sm.release("foo.ipynb")
    assert ok is True
    # OpenCode delete called for the active one (sid-B).
    assert any(
        c[0] == "delete_session" and c[1] == "sid-B" for c in client.calls
    )
    # Both bindings cleared.
    assert sm.list_sessions_for_notebook("foo.ipynb") == []
    assert sm.get_active("foo.ipynb") is None


@pytest.mark.asyncio
async def test_get_or_create_after_set_active_returns_the_active_one():
    client = FakeClient()
    sm = SessionManager(lambda: client)
    first = await sm.get_or_create("foo.ipynb")
    sm.bind_existing("foo.ipynb", "another-sid")
    sm.set_active("foo.ipynb", "another-sid")
    # Next /edit-style call returns the manually-set active session.
    assert await sm.get_or_create("foo.ipynb") == "another-sid"
    # No additional create_session call.
    creates = [c for c in client.calls if c[0] == "create_session"]
    assert len(creates) == 1  # only the initial get_or_create