test_browser.py

"""Tests for browser.py multi-session support."""

import asyncio
import json
from unittest.mock import MagicMock

import pytest
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer

from browser import BrowserClient, build_app


# ── Unit tests: BrowserClient session registry ──────────────────────


class TestSessionRegistry:
    """Test register/unregister/default promotion logic."""

    def test_register_session(self):
        c = BrowserClient()
        c.register_session("sess-1", "tab-1")
        assert c.sessions == {"sess-1": "tab-1"}
        assert c._default_session == "sess-1"

    def test_register_multiple_sessions(self):
        c = BrowserClient()
        c.register_session("sess-1", "tab-1")
        c.register_session("sess-2", "tab-2")
        assert len(c.sessions) == 2
        assert c._default_session == "sess-2"

    def test_unregister_non_default(self):
        c = BrowserClient()
        c.register_session("sess-1", "tab-1")
        c.register_session("sess-2", "tab-2")
        c.unregister_session("sess-1")
        assert "sess-1" not in c.sessions
        assert c._default_session == "sess-2"

    def test_unregister_default_promotes_next(self):
        c = BrowserClient()
        c.register_session("sess-1", "tab-1")
        c.register_session("sess-2", "tab-2")
        c.unregister_session("sess-2")
        assert c._default_session == "sess-1"

    def test_unregister_last_session(self):
        c = BrowserClient()
        c.register_session("sess-1", "tab-1")
        c.unregister_session("sess-1")
        assert c.sessions == {}
        assert c._default_session is None

    def test_unregister_nonexistent(self):
        c = BrowserClient()
        c.register_session("sess-1", "tab-1")
        c.unregister_session("bogus")
        assert c.sessions == {"sess-1": "tab-1"}

    def test_get_target_id(self):
        c = BrowserClient()
        c.register_session("sess-1", "tab-1")
        assert c.get_target_id("sess-1") == "tab-1"
        assert c.get_target_id("bogus") is None

    def test_backwards_compat_properties(self):
        c = BrowserClient()
        assert c.session_id is None
        assert c.target_id is None
        c.register_session("sess-1", "tab-1")
        assert c.session_id == "sess-1"
        assert c.target_id == "tab-1"

    def test_auto_cleanup_on_target_destroyed(self):
        """Simulate _read_loop receiving Target.targetDestroyed event."""
        c = BrowserClient()
        c.register_session("sess-1", "tab-1")
        c.register_session("sess-2", "tab-2")
        # Simulate the logic from _read_loop
        destroyed_tid = "tab-1"
        for sid, tid in list(c.sessions.items()):
            if tid == destroyed_tid:
                c.unregister_session(sid)
        assert "sess-1" not in c.sessions
        assert "sess-2" in c.sessions
        assert c._default_session == "sess-2"


# ── Integration tests: HTTP API with mocked CDP ─────────────────────


def make_mock_client():
    """Create a BrowserClient with mocked CDP/evaluate methods."""
    client = BrowserClient()
    client.ws = MagicMock()
    client.ws.closed = False

    async def mock_cdp(method, params=None, session_id=None):
        if method == "Target.getTargets":
            return {"result": {"targetInfos": [
                {"targetId": "tab-1", "type": "page", "title": "Page One", "url": "https://one.com"},
                {"targetId": "tab-2", "type": "page", "title": "Page Two", "url": "https://two.com"},
            ]}}
        if method == "Target.createTarget":
            return {"result": {"targetId": "tab-new"}}
        if method == "Target.attachToTarget":
            tid = (params or {}).get("targetId", "unknown")
            return {"result": {"sessionId": f"session-{tid}-123"}}
        if method == "Target.detachFromTarget":
            return {"result": {}}
        if method == "Target.closeTarget":
            return {"result": {"success": True}}
        if method == "Runtime.evaluate":
            expr = (params or {}).get("expression", "")
            if "document.title" in expr:
                tid = client.sessions.get(session_id, "?")
                return {"result": {"result": {"value": f"Title for {tid}", "type": "string"}}}
            if "document.body.innerText" in expr:
                tid = client.sessions.get(session_id, "?")
                return {"result": {"result": {"value": f"Text for {tid}", "type": "string"}}}
            if "window.location.href" in expr:
                tid = client.sessions.get(session_id, "?")
                return {"result": {"result": {"value": f"https://{tid}.example.com", "type": "string"}}}
            if "readyState" in expr:
                return {"result": {"result": {"value": "complete", "type": "string"}}}
            if "querySelector" in expr:
                return {"result": {"result": {"value": "clicked", "type": "string"}}}
            return {"result": {"result": {"value": None, "type": "undefined"}}}
        if method == "Page.captureScreenshot":
            return {"result": {"data": "iVBORw=="}}
        if method == "Page.navigate":
            return {"result": {"frameId": "frame-1"}}
        return {"result": {}}

    client.cdp = mock_cdp
    return client


@pytest.fixture
def mock_client():
    return make_mock_client()


async def make_http(mock_client):
    """Create a test HTTP client from a mock BrowserClient."""
    app = build_app(mock_client)
    server = TestServer(app)
    client = TestClient(server)
    await client.start_server()
    return client


# ── /status ──────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_status_empty(mock_client):
    http = await make_http(mock_client)
    try:
        resp = await http.get("/status")
        data = await resp.json()
        assert data["connected"] is True
        assert data["sessions"] == {}
        assert data["default_session"] is None
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_status_with_sessions(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    mock_client.register_session("sess-B", "tab-B")
    http = await make_http(mock_client)
    try:
        resp = await http.get("/status")
        data = await resp.json()
        assert len(data["sessions"]) == 2
        assert data["default_session"] == "sess-B"
        assert data["session_id"] == "sess-B"
        assert data["target_id"] == "tab-B"
    finally:
        await http.close()


# ── /tabs ────────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_tabs(mock_client):
    http = await make_http(mock_client)
    try:
        resp = await http.get("/tabs")
        data = await resp.json()
        assert len(data) == 2
        assert data[0]["id"] == "tab-1"
        assert data[1]["url"] == "https://two.com"
    finally:
        await http.close()


# ── /attach ──────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_attach_registers_session(mock_client):
    http = await make_http(mock_client)
    try:
        resp = await http.post("/attach", json={"target_id": "tab-1"})
        data = await resp.json()
        assert data["session_id"] == "session-tab-1-123"
        assert "session-tab-1-123" in mock_client.sessions
        assert mock_client._default_session == "session-tab-1-123"
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_attach_two_tabs(mock_client):
    http = await make_http(mock_client)
    try:
        r1 = await http.post("/attach", json={"target_id": "tab-1"})
        r2 = await http.post("/attach", json={"target_id": "tab-2"})
        d1 = await r1.json()
        d2 = await r2.json()
        assert d1["session_id"] != d2["session_id"]
        assert len(mock_client.sessions) == 2
        assert mock_client._default_session == d2["session_id"]
    finally:
        await http.close()


# ── /navigate ────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_navigate_creates_new_tab(mock_client):
    http = await make_http(mock_client)
    try:
        resp = await http.post("/navigate", json={"url": "https://test.com"})
        data = await resp.json()
        assert data["target_id"] == "tab-new"
        assert data["session_id"] in mock_client.sessions
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_navigate_existing_session(mock_client):
    http = await make_http(mock_client)
    try:
        r1 = await http.post("/navigate", json={"url": "https://one.com"})
        sid = (await r1.json())["session_id"]
        r2 = await http.post("/navigate", json={"url": "https://two.com", "session": sid, "new_tab": False})
        data = await r2.json()
        assert data["session_id"] == sid
        assert len(mock_client.sessions) == 1
    finally:
        await http.close()


# ── /eval with session isolation ─────────────────────────────────────


@pytest.mark.asyncio
async def test_eval_default_session(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/eval", json={"js": "document.title"})
        data = await resp.json()
        assert data["value"] == "Title for tab-A"
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_eval_explicit_session(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    mock_client.register_session("sess-B", "tab-B")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/eval", json={"js": "document.title", "session": "sess-A"})
        data = await resp.json()
        assert data["value"] == "Title for tab-A"
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_eval_no_session_returns_error(mock_client):
    http = await make_http(mock_client)
    try:
        resp = await http.post("/eval", json={"js": "document.title"})
        assert resp.status == 400
        data = await resp.json()
        assert "error" in data
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_eval_unknown_session_returns_error(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/eval", json={"js": "document.title", "session": "bogus"})
        assert resp.status == 400
        data = await resp.json()
        assert "Unknown session" in data["error"]
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_eval_isolation(mock_client):
    """Two sessions each get their own tab's result."""
    mock_client.register_session("sess-A", "tab-A")
    mock_client.register_session("sess-B", "tab-B")
    http = await make_http(mock_client)
    try:
        rA = await http.post("/eval", json={"js": "document.title", "session": "sess-A"})
        rB = await http.post("/eval", json={"js": "document.title", "session": "sess-B"})
        assert (await rA.json())["value"] == "Title for tab-A"
        assert (await rB.json())["value"] == "Title for tab-B"
    finally:
        await http.close()


# ── GET endpoints with ?session= ─────────────────────────────────────


@pytest.mark.asyncio
async def test_title_query_param(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    mock_client.register_session("sess-B", "tab-B")
    http = await make_http(mock_client)
    try:
        resp = await http.get("/title?session=sess-A")
        data = await resp.json()
        assert data["value"] == "Title for tab-A"
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_text_query_param(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.get("/text?session=sess-A")
        data = await resp.json()
        assert data["value"] == "Text for tab-A"
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_url_query_param(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.get("/url?session=sess-A")
        data = await resp.json()
        assert data["value"] == "https://tab-A.example.com"
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_get_no_session_error(mock_client):
    http = await make_http(mock_client)
    try:
        for endpoint in ["/title", "/text", "/url", "/screenshot"]:
            resp = await http.get(endpoint)
            assert resp.status == 400, f"{endpoint} should 400 with no session"
    finally:
        await http.close()


# ── /click ───────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_click_with_session(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/click", json={"selector": "#btn", "session": "sess-A"})
        data = await resp.json()
        assert data["value"] == "clicked"
    finally:
        await http.close()


# ── /detach ──────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_detach_specific_session(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    mock_client.register_session("sess-B", "tab-B")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/detach", json={"session": "sess-A"})
        data = await resp.json()
        assert data["detached"] == "tab-A"
        assert "sess-A" not in mock_client.sessions
        assert "sess-B" in mock_client.sessions
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_detach_default_promotes_other(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    mock_client.register_session("sess-B", "tab-B")
    http = await make_http(mock_client)
    try:
        await http.post("/detach", json={"session": "sess-B"})
        assert mock_client._default_session == "sess-A"
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_detach_no_body_uses_default(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/detach")
        data = await resp.json()
        assert data["detached"] == "tab-A"
        assert mock_client.sessions == {}
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_detach_unknown_session_error(mock_client):
    http = await make_http(mock_client)
    try:
        resp = await http.post("/detach", json={"session": "bogus"})
        assert resp.status == 400
    finally:
        await http.close()


# ── /close ───────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_close_by_target_id(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/close", json={"target_id": "tab-A"})
        data = await resp.json()
        assert data.get("success") is True
        assert "sess-A" not in mock_client.sessions
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_close_by_session(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/close", json={"session": "sess-A"})
        data = await resp.json()
        assert data.get("success") is True
        assert mock_client.sessions == {}
    finally:
        await http.close()


# ── /cdp raw ─────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_cdp_raw_with_session(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.post("/cdp", json={
            "method": "Runtime.evaluate",
            "params": {"expression": "1+1"},
            "session": "sess-A",
        })
        assert resp.status == 200
    finally:
        await http.close()


@pytest.mark.asyncio
async def test_cdp_raw_session_false(mock_client):
    http = await make_http(mock_client)
    try:
        resp = await http.post("/cdp", json={
            "method": "Target.getTargets",
            "session": False,
        })
        assert resp.status == 200
    finally:
        await http.close()


# ── /screenshot ──────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_screenshot_with_session(mock_client):
    mock_client.register_session("sess-A", "tab-A")
    http = await make_http(mock_client)
    try:
        resp = await http.get("/screenshot?session=sess-A")
        data = await resp.json()
        assert "data_length" in data
    finally:
        await http.close()