snakebridge

"""Tests for SnakeBridge gRPC streaming client behavior."""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import snakepit_bridge_pb2 as pb2
from google.protobuf import any_pb2

from snakebridge_client import BridgeClient


class _FakeStream:
    def __init__(self, chunks):
        self._it = iter(chunks)
        self.cancelled = False

    def __iter__(self):
        return self

    def __next__(self):
        return next(self._it)

    def cancel(self):
        self.cancelled = True


class _FakeStub:
    def __init__(self):
        self.streaming_calls = []
        self.execute_calls = []

    def ExecuteStreamingTool(self, request, metadata=None, timeout=None):
        self.streaming_calls.append((request, metadata, timeout))
        return _FakeStream(
            [
                pb2.ToolChunk(
                    chunk_id="c1",
                    data=b'{"step":1,"total":2}',
                    is_final=False,
                    metadata={"k": "v"},
                ),
                pb2.ToolChunk(
                    chunk_id="c2",
                    data=b'{"status":"done"}',
                    is_final=True,
                    metadata={"execution_time_ms": "5"},
                ),
            ]
        )

    def ExecuteTool(self, *args, **kwargs):
        self.execute_calls.append((args, kwargs))
        raise AssertionError("ExecuteTool should not be used for streaming")


class _UnaryStub:
    def __init__(self):
        self.execute_calls = []

    def ExecuteTool(self, request, metadata=None, timeout=None):
        self.execute_calls.append((request, metadata, timeout))
        return pb2.ExecuteToolResponse(
            success=True,
            result=any_pb2.Any(
                type_url="type.googleapis.com/google.protobuf.StringValue",
                value=b'"ok"',
            ),
        )


def test_execute_tool_sets_request_metadata_and_correlation_header():
    stub = _UnaryStub()
    client = BridgeClient(stub=stub)

    result = client.execute_tool(
        "session-1",
        "tool",
        {"answer": 42},
        request_metadata={"thread_sensitive": True, "extra": "1"},
        correlation_id="cid-123",
    )

    assert result == "ok"
    assert len(stub.execute_calls) == 1

    request, call_md, _timeout = stub.execute_calls[0]
    assert request.session_id == "session-1"
    assert request.tool_name == "tool"
    assert request.parameters["answer"].value == b"42"
    assert request.metadata["correlation_id"] == "cid-123"
    assert request.metadata["thread_sensitive"] == "True"
    assert request.metadata["extra"] == "1"

    assert any(
        k == "x-snakepit-correlation-id" and v == "cid-123" for (k, v) in (call_md or [])
    )


def test_execute_tool_rejects_non_bytes_binary_parameters():
    stub = _UnaryStub()
    client = BridgeClient(stub=stub)

    try:
        client.execute_tool(
            "s",
            "t",
            binary_parameters={"bad": "not-bytes"},
            correlation_id="cid",
        )
    except TypeError as exc:
        assert "binary_parameters" in str(exc)
    else:
        raise AssertionError("Expected TypeError for non-bytes binary_parameters")


def test_execute_streaming_tool_uses_streaming_rpc_and_sets_correlation_header():
    stub = _FakeStub()
    client = BridgeClient(stub=stub)

    out = list(
        client.execute_streaming_tool(
            "session-1",
            "stream_tool",
            {"steps": 2},
            binary_parameters={"blob": b"\x00\x01"},
            correlation_id="cid-123",
        )
    )

    assert len(stub.streaming_calls) == 1
    assert stub.execute_calls == []
    request, call_md, _timeout = stub.streaming_calls[0]

    # Ensures we used ExecuteStreamingTool and encoded parameters in Any.value as raw JSON.
    assert request.session_id == "session-1"
    assert request.tool_name == "stream_tool"
    assert "steps" in request.parameters
    assert request.parameters["steps"].type_url == "type.googleapis.com/google.protobuf.StringValue"
    assert request.parameters["steps"].value == b"2"
    assert request.binary_parameters["blob"] == b"\x00\x01"

    # Correlation header is present on RPC metadata.
    assert any(k == "x-snakepit-correlation-id" and v == "cid-123" for (k, v) in (call_md or []))

    # Chunk decoding behavior.
    assert out[0]["step"] == 1
    assert out[0]["total"] == 2
    assert out[0]["is_final"] is False
    assert out[0]["_metadata"]["k"] == "v"
    assert out[0]["_chunk_id"] == "c1"

    assert out[-1]["status"] == "done"
    assert out[-1]["is_final"] is True
    assert out[-1]["_metadata"]["execution_time_ms"] == "5"
    assert out[-1]["_chunk_id"] == "c2"


def test_execute_streaming_tool_with_empty_and_non_json_chunks():
    """Test handling of chunks with empty or non-JSON data."""

    class EmptyChunkStub:
        def __init__(self):
            self.calls = []

        def ExecuteStreamingTool(self, request, metadata=None, timeout=None):
            self.calls.append((request, metadata, timeout))
            return _FakeStream(
                [
                    pb2.ToolChunk(chunk_id="e1", data=b"", is_final=False),
                    pb2.ToolChunk(chunk_id="e2", data=b"not json", is_final=True),
                ]
            )

    stub = EmptyChunkStub()
    client = BridgeClient(stub=stub)

    out = list(client.execute_streaming_tool("s", "t", correlation_id="cid"))

    # Empty data returns empty dict + is_final
    assert out[0] == {"is_final": False, "_chunk_id": "e1"}

    # Non-JSON data is preserved as bytes in "data" key
    assert out[1]["data"] == b"not json"
    assert out[1]["is_final"] is True
    assert out[1]["_chunk_id"] == "e2"


def test_execute_streaming_tool_auto_generates_correlation_id():
    """Test that correlation ID is auto-generated if not provided."""
    stub = _FakeStub()
    client = BridgeClient(stub=stub)

    # Don't pass correlation_id
    list(client.execute_streaming_tool("s", "t"))

    assert len(stub.streaming_calls) == 1
    _request, call_md, _timeout = stub.streaming_calls[0]

    # Should have generated a correlation ID
    corr_headers = [v for (k, v) in (call_md or []) if k == "x-snakepit-correlation-id"]
    assert len(corr_headers) == 1
    assert corr_headers[0]


def test_execute_streaming_tool_stops_on_final_chunk():
    class ExtraChunkStub:
        def __init__(self):
            self.calls = []

        def ExecuteStreamingTool(self, request, metadata=None, timeout=None):
            self.calls.append((request, metadata, timeout))
            return _FakeStream(
                [
                    pb2.ToolChunk(chunk_id="c1", data=b'{"progress":1}', is_final=False),
                    pb2.ToolChunk(chunk_id="c2", data=b'{"progress":2}', is_final=True),
                    pb2.ToolChunk(chunk_id="c3", data=b'{"progress":3}', is_final=False),
                ]
            )

    stub = ExtraChunkStub()
    client = BridgeClient(stub=stub)

    out = list(client.execute_streaming_tool("s", "t", correlation_id="cid"))

    assert [chunk.get("progress") for chunk in out] == [1, 2]


def test_execute_streaming_tool_rejects_non_bytes_binary_parameters():
    stub = _FakeStub()
    client = BridgeClient(stub=stub)

    try:
        list(
            client.execute_streaming_tool(
                "s",
                "t",
                binary_parameters={"bad": "not-bytes"},
                correlation_id="cid",
            )
        )
    except TypeError as exc:
        assert "binary_parameters" in str(exc)
    else:
        raise AssertionError("Expected TypeError for non-bytes binary_parameters")