"""Tests for LangGraph stream payload helpers."""
import pytest

from services.stream_payloads import normalize_custom_writer_payload, parse_stream_graph_item


@pytest.mark.parametrize(
    "raw,expected_kind,extra_keys",
    [
        ("hello", "tool", {"message": "hello"}),
        ({"kind": "tool", "message": "x"}, "tool", {"message": "x"}),
        ({"foo": 1}, "custom", {"foo": 1}),
        ({"kind": "", "x": 1}, "custom", {"x": 1}),
    ],
)
def test_normalize_custom_writer_payload(raw, expected_kind, extra_keys):
    out = normalize_custom_writer_payload(raw)
    assert out["kind"] == expected_kind
    for k, v in extra_keys.items():
        assert out[k] == v


def test_parse_stream_graph_item_two_tuple():
    assert parse_stream_graph_item(("custom", "x")) == ("custom", "x")
    assert parse_stream_graph_item(("messages", {"a": 1})) == ("messages", {"a": 1})


def test_parse_stream_graph_item_three_tuple():
    assert parse_stream_graph_item(((), "custom", {"k": 1})) == ("custom", {"k": 1})


def test_parse_stream_graph_item_dict_custom():
    assert parse_stream_graph_item({"type": "custom", "data": {"z": 2}}) == ("custom", {"z": 2})