python-sdk/tests/test_filesystem_pty.py

from __future__ import annotations

import base64
import json
from unittest.mock import AsyncMock, MagicMock

import pytest
import respx

from wrenn.capsule import Capsule
from wrenn.models import FileEntry
from wrenn.pty import (
    AsyncPtySession,
    PtyEventType,
    PtySession,
    _parse_pty_event,
)

BASE = "https://app.wrenn.dev/api"


def _make_capsule(cap_id: str = "cl-abc") -> Capsule:
    respx.post(f"{BASE}/v1/capsules").respond(
        201, json={"id": cap_id, "status": "running"}
    )
    return Capsule(api_key="wrn_test1234567890abcdef12345678", base_url=BASE)


class TestFilesRead:
    @respx.mock
    def test_read_returns_string(self):
        cap = _make_capsule()
        content = b"file contents here"
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/read").respond(
            200, content=content
        )
        data = cap.files.read("/app/main.py")
        assert data == "file contents here"

    @respx.mock
    def test_read_bytes(self):
        cap = _make_capsule()
        content = b"\x00\x01\x02"
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/read").respond(
            200, content=content
        )
        data = cap.files.read_bytes("/bin/binary")
        assert data == b"\x00\x01\x02"


class TestFilesWrite:
    @respx.mock
    def test_write_string(self):
        cap = _make_capsule()
        route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/write").respond(204)
        cap.files.write("/app/main.py", "print('hello')")
        assert route.called

    @respx.mock
    def test_write_bytes(self):
        cap = _make_capsule()
        route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/write").respond(204)
        cap.files.write("/app/data.bin", b"\x00\x01\x02")
        assert route.called


class TestFilesList:
    @respx.mock
    def test_list_returns_entries(self):
        cap = _make_capsule()
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
            200,
            json={
                "entries": [
                    {
                        "name": "main.py",
                        "path": "/home/wrenn-user/main.py",
                        "type": "file",
                        "size": 1024,
                        "mode": 33188,
                        "permissions": "-rw-r--r--",
                        "owner": "wrenn-user",
                        "group": "wrenn-user",
                        "modified_at": 1712899200,
                        "symlink_target": None,
                    },
                    {
                        "name": "config",
                        "path": "/home/wrenn-user/config",
                        "type": "directory",
                        "size": 4096,
                        "mode": 16877,
                        "permissions": "drwxr-xr-x",
                        "owner": "wrenn-user",
                        "group": "wrenn-user",
                        "modified_at": 1712899100,
                        "symlink_target": None,
                    },
                ]
            },
        )
        entries = cap.files.list("/home/wrenn-user")
        assert len(entries) == 2
        assert isinstance(entries[0], FileEntry)
        assert entries[0].name == "main.py"
        assert entries[0].type == "file"
        assert entries[1].name == "config"
        assert entries[1].type == "directory"

    @respx.mock
    def test_list_with_depth(self):
        cap = _make_capsule()
        route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
            200, json={"entries": []}
        )
        cap.files.list("/home/wrenn-user", depth=3)
        body = json.loads(route.calls[0].request.content)
        assert body["depth"] == 3

    @respx.mock
    def test_list_empty(self):
        cap = _make_capsule()
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
            200, json={"entries": []}
        )
        entries = cap.files.list("/empty")
        assert entries == []


class TestFilesMakeDir:
    @respx.mock
    def test_make_dir_returns_entry(self):
        cap = _make_capsule()
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/mkdir").respond(
            200,
            json={
                "entry": {
                    "name": "data",
                    "path": "/home/wrenn-user/data",
                    "type": "directory",
                    "size": 4096,
                    "mode": 16877,
                    "permissions": "drwxr-xr-x",
                    "owner": "wrenn-user",
                    "group": "wrenn-user",
                    "modified_at": 1712899200,
                    "symlink_target": None,
                }
            },
        )
        entry = cap.files.make_dir("/home/wrenn-user/data")
        assert isinstance(entry, FileEntry)
        assert entry.name == "data"
        assert entry.type == "directory"

    @respx.mock
    def test_make_dir_existing_returns_gracefully(self):
        cap = _make_capsule()
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/mkdir").respond(
            409,
            json={"error": {"code": "conflict", "message": "already exists"}},
        )
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
            200,
            json={
                "entries": [
                    {
                        "name": "data",
                        "path": "/home/wrenn-user/data",
                        "type": "directory",
                        "size": 4096,
                        "mode": 16877,
                        "permissions": "drwxr-xr-x",
                        "owner": "wrenn-user",
                        "group": "wrenn-user",
                        "modified_at": 1712899200,
                        "symlink_target": None,
                    }
                ]
            },
        )
        entry = cap.files.make_dir("/home/wrenn-user/data")
        assert entry.name == "data"


class TestFilesRemove:
    @respx.mock
    def test_remove_succeeds(self):
        cap = _make_capsule()
        route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/remove").respond(204)
        cap.files.remove("/home/wrenn-user/old_data")
        assert route.called

    @respx.mock
    def test_remove_sends_path(self):
        cap = _make_capsule()
        route = respx.post(f"{BASE}/v1/capsules/cl-abc/files/remove").respond(204)
        cap.files.remove("/tmp/test.txt")
        body = json.loads(route.calls[0].request.content)
        assert body["path"] == "/tmp/test.txt"


class TestFilesExists:
    @respx.mock
    def test_exists_true(self):
        cap = _make_capsule()
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
            200,
            json={
                "entries": [
                    {"name": "hello.txt", "path": "/tmp/hello.txt", "type": "file"}
                ]
            },
        )
        assert cap.files.exists("/tmp/hello.txt") is True

    @respx.mock
    def test_exists_false(self):
        cap = _make_capsule()
        respx.post(f"{BASE}/v1/capsules/cl-abc/files/list").respond(
            200, json={"entries": []}
        )
        assert cap.files.exists("/tmp/nope.txt") is False


class TestPtyEventParsing:
    def test_started_event(self):
        raw = {"type": "started", "tag": "pty-a1b2c3d4", "pid": 42}
        event = _parse_pty_event(raw)
        assert event.type == PtyEventType.started
        assert event.pid == 42
        assert event.tag == "pty-a1b2c3d4"

    def test_output_event_base64(self):
        encoded = base64.b64encode(b"ls -la\n").decode()
        raw = {"type": "output", "data": encoded}
        event = _parse_pty_event(raw)
        assert event.type == PtyEventType.output
        assert event.data == b"ls -la\n"

    def test_output_event_empty(self):
        raw = {"type": "output", "data": ""}
        event = _parse_pty_event(raw)
        assert event.data == b""

    def test_exit_event(self):
        raw = {"type": "exit", "exit_code": 0}
        event = _parse_pty_event(raw)
        assert event.type == PtyEventType.exit
        assert event.exit_code == 0

    def test_error_event(self):
        raw = {"type": "error", "data": "process not found", "fatal": True}
        event = _parse_pty_event(raw)
        assert event.type == PtyEventType.error
        assert event.data == "process not found"
        assert event.fatal is True

    def test_ping_event(self):
        raw = {"type": "ping"}
        event = _parse_pty_event(raw)
        assert event.type == PtyEventType.ping


class TestPtySessionWrite:
    def test_write_sends_base64_input(self):
        ws = MagicMock()
        session = PtySession(ws, "cl-abc")
        session.write(b"ls -la\n")
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "input"
        assert base64.b64decode(sent["data"]) == b"ls -la\n"


class TestPtySessionResize:
    def test_resize_sends_dimensions(self):
        ws = MagicMock()
        session = PtySession(ws, "cl-abc")
        session.resize(120, 40)
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "resize"
        assert sent["cols"] == 120
        assert sent["rows"] == 40

    def test_resize_zero_raises(self):
        ws = MagicMock()
        session = PtySession(ws, "cl-abc")
        with pytest.raises(ValueError, match="greater than 0"):
            session.resize(0, 40)
        with pytest.raises(ValueError, match="greater than 0"):
            session.resize(80, 0)


class TestPtySessionKill:
    def test_kill_sends_message(self):
        ws = MagicMock()
        session = PtySession(ws, "cl-abc")
        session.kill()
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "kill"


class TestPtySessionIteration:
    def test_iter_yields_events_until_exit(self):
        ws = MagicMock()
        messages = [
            json.dumps({"type": "started", "tag": "pty-abc12345", "pid": 1}),
            json.dumps({"type": "output", "data": base64.b64encode(b"hello").decode()}),
            json.dumps({"type": "exit", "exit_code": 0}),
        ]
        ws.receive_text.side_effect = messages
        session = PtySession(ws, "cl-abc")
        events = list(session)
        assert len(events) == 3
        assert events[0].type == PtyEventType.started
        assert session.tag == "pty-abc12345"
        assert session.pid == 1
        assert events[1].type == PtyEventType.output
        assert events[1].data == b"hello"
        assert events[2].type == PtyEventType.exit
        assert events[2].exit_code == 0

    def test_iter_stops_on_fatal_error(self):
        ws = MagicMock()
        messages = [
            json.dumps({"type": "error", "data": "fatal", "fatal": True}),
        ]
        ws.receive_text.side_effect = messages
        session = PtySession(ws, "cl-abc")
        events = list(session)
        assert len(events) == 1
        assert events[0].type == PtyEventType.error

    def test_iter_stops_on_disconnect(self):
        import httpx_ws

        ws = MagicMock()
        ws.receive_text.side_effect = httpx_ws.WebSocketDisconnect()
        session = PtySession(ws, "cl-abc")
        events = list(session)
        assert events == []


class TestPtySessionPong:
    def test_ping_triggers_pong(self):
        ws = MagicMock()
        ws.receive_text.side_effect = [
            json.dumps({"type": "ping"}),
            json.dumps({"type": "exit", "exit_code": 0}),
        ]
        session = PtySession(ws, "cl-abc")
        events = list(session)
        assert events[0].type == PtyEventType.ping
        sent = [json.loads(c[0][0]) for c in ws.send_text.call_args_list]
        assert {"type": "pong"} in sent

    def test_no_pong_without_ping(self):
        ws = MagicMock()
        ws.receive_text.side_effect = [
            json.dumps({"type": "output", "data": ""}),
            json.dumps({"type": "exit", "exit_code": 0}),
        ]
        session = PtySession(ws, "cl-abc")
        list(session)
        sent = [json.loads(c[0][0]) for c in ws.send_text.call_args_list]
        assert {"type": "pong"} not in sent

    def test_send_pong_swallows_closed_ws(self):
        import httpx_ws

        ws = MagicMock()
        ws.send_text.side_effect = httpx_ws.WebSocketNetworkError()
        session = PtySession(ws, "cl-abc")
        session._send_pong()  # must not raise


class TestPtySessionContextManager:
    def test_exit_kills_and_closes(self):
        ws = MagicMock()
        session = PtySession(ws, "cl-abc")
        with session:
            pass
        ws.send_text.assert_called()
        ws.close.assert_called()

    def test_exit_ignores_errors(self):
        ws = MagicMock()
        ws.send_text.side_effect = Exception("already closed")
        session = PtySession(ws, "cl-abc")
        with session:
            pass


class TestPtySessionSendStart:
    def test_send_start_with_defaults(self):
        ws = MagicMock()
        session = PtySession(ws, "cl-abc")
        session._send_start()
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "start"
        assert sent["cmd"] == "/bin/bash"
        assert sent["cols"] == 80
        assert sent["rows"] == 24

    def test_send_start_with_all_params(self):
        ws = MagicMock()
        session = PtySession(ws, "cl-abc")
        session._send_start(
            cmd="/bin/zsh",
            args=["-l"],
            cols=120,
            rows=40,
            envs={"TERM": "xterm-256color"},
            cwd="/home/wrenn-user",
        )
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["cmd"] == "/bin/zsh"
        assert sent["args"] == ["-l"]
        assert sent["cols"] == 120


class TestPtySessionSendConnect:
    def test_send_connect(self):
        ws = MagicMock()
        session = PtySession(ws, "cl-abc")
        session._send_connect("pty-abc12345")
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "connect"
        assert sent["tag"] == "pty-abc12345"


class TestAsyncPtySession:
    @pytest.mark.asyncio
    async def test_async_write_sends_base64(self):
        ws = AsyncMock()
        session = AsyncPtySession(ws, "cl-abc")
        await session.write(b"hello")
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "input"
        assert base64.b64decode(sent["data"]) == b"hello"

    @pytest.mark.asyncio
    async def test_async_resize(self):
        ws = AsyncMock()
        session = AsyncPtySession(ws, "cl-abc")
        await session.resize(100, 30)
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "resize"
        assert sent["cols"] == 100
        assert sent["rows"] == 30

    @pytest.mark.asyncio
    async def test_async_resize_zero_raises(self):
        ws = AsyncMock()
        session = AsyncPtySession(ws, "cl-abc")
        with pytest.raises(ValueError):
            await session.resize(0, 10)

    @pytest.mark.asyncio
    async def test_async_kill(self):
        ws = AsyncMock()
        session = AsyncPtySession(ws, "cl-abc")
        await session.kill()
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "kill"

    @pytest.mark.asyncio
    async def test_async_context_manager(self):
        ws = AsyncMock()
        session = AsyncPtySession(ws, "cl-abc")
        async with session:
            pass
        ws.send_text.assert_called()
        ws.close.assert_called()

    @pytest.mark.asyncio
    async def test_async_send_start(self):
        ws = AsyncMock()
        session = AsyncPtySession(ws, "cl-abc")
        await session._send_start(cmd="/bin/zsh", cols=100, rows=30)
        sent = json.loads(ws.send_text.call_args[0][0])
        assert sent["type"] == "start"
        assert sent["cmd"] == "/bin/zsh"
        assert sent["cols"] == 100

    @pytest.mark.asyncio
    async def test_async_ping_triggers_pong(self):
        ws = AsyncMock()
        ws.receive_text.side_effect = [
            json.dumps({"type": "ping"}),
            json.dumps({"type": "exit", "exit_code": 0}),
        ]
        session = AsyncPtySession(ws, "cl-abc")
        events = [e async for e in session]
        assert events[0].type == PtyEventType.ping
        sent = [json.loads(c[0][0]) for c in ws.send_text.call_args_list]
        assert {"type": "pong"} in sent

    @pytest.mark.asyncio
    async def test_async_send_pong_swallows_closed_ws(self):
        import httpx_ws

        ws = AsyncMock()
        ws.send_text.side_effect = httpx_ws.WebSocketNetworkError()
        session = AsyncPtySession(ws, "cl-abc")
        await session._send_pong()  # must not raise

    @pytest.mark.asyncio
    async def test_async_iteration(self):
        ws = AsyncMock()
        messages = [
            json.dumps({"type": "started", "tag": "pty-xyz", "pid": 5}),
            json.dumps({"type": "output", "data": base64.b64encode(b"hi").decode()}),
            json.dumps({"type": "exit", "exit_code": 0}),
        ]
        ws.receive_text.side_effect = messages
        session = AsyncPtySession(ws, "cl-abc")
        events = []
        async for event in session:
            events.append(event)
        assert len(events) == 3
        assert events[0].type == PtyEventType.started
        assert session.tag == "pty-xyz"
        assert session.pid == 5
        assert events[2].type == PtyEventType.exit


class TestExports:
    def test_file_entry_importable(self):
        from wrenn import FileEntry as FE

        assert FE is not None

    def test_pty_session_importable(self):
        from wrenn import PtySession as PS

        assert PS is not None

    def test_async_pty_session_importable(self):
        from wrenn import AsyncPtySession as APS

        assert APS is not None

    def test_pty_event_importable(self):
        from wrenn import PtyEvent as PE
        from wrenn import PtyEventType as PET

        assert PE is not None
        assert PET is not None