nordisk/knut

python/tests/test_proto.py

"""Tests for the pure-Python proto encoder (no server required)."""

from knut._proto import (
    COMMAND_STATEMENT_QUERY,
    COMMAND_STATEMENT_UPDATE,
    encode_command,
    decode_ticket,
    _varint,
    _bytes_field,
    _string_field,
    _parse_fields,
)


class TestVarint:
    def test_single_byte(self):
        assert _varint(0) == b"\x00"
        assert _varint(1) == b"\x01"
        assert _varint(127) == b"\x7f"

    def test_multi_byte(self):
        assert _varint(128) == b"\x80\x01"
        assert _varint(300) == b"\xac\x02"


class TestFieldEncoding:
    def test_bytes_field_tag(self):
        # Field 1, wire type 2 → tag byte = (1<<3)|2 = 0x0a
        encoded = _bytes_field(1, b"hello")
        assert encoded[0] == 0x0A
        assert encoded[1] == 5           # length varint
        assert encoded[2:] == b"hello"

    def test_string_field_roundtrip(self):
        fields = _parse_fields(_string_field(1, "hello world"))
        assert fields[1] == b"hello world"

    def test_field_2(self):
        # Field 2, wire type 2 → tag byte = (2<<3)|2 = 0x12
        encoded = _bytes_field(2, b"x")
        assert encoded[0] == 0x12


class TestEncodeCommand:
    def _parse_any(self, buf: bytes) -> tuple[str, bytes]:
        """Parse an Any message → (type_url, value_bytes)."""
        fields = _parse_fields(buf)
        return fields[1].decode(), fields[2]

    def test_type_url_present(self):
        buf = encode_command(COMMAND_STATEMENT_QUERY, "SELECT 1")
        type_url, _ = self._parse_any(buf)
        assert type_url == COMMAND_STATEMENT_QUERY

    def test_query_preserved(self):
        query = "MATCH (n:Person) RETURN n.id"
        buf = encode_command(COMMAND_STATEMENT_QUERY, query)
        _, inner = self._parse_any(buf)
        inner_fields = _parse_fields(inner)
        assert inner_fields[1].decode() == query

    def test_update_type_url(self):
        buf = encode_command(COMMAND_STATEMENT_UPDATE, "VERTICES:Person")
        type_url, _ = self._parse_any(buf)
        assert type_url == COMMAND_STATEMENT_UPDATE

    def test_empty_query(self):
        buf = encode_command(COMMAND_STATEMENT_QUERY, "")
        _, inner = self._parse_any(buf)
        inner_fields = _parse_fields(inner)
        assert inner_fields.get(1, b"") == b""

    def test_unicode_query(self):
        query = "MATCH (n) WHERE n.name = 'Ångström' RETURN n"
        buf = encode_command(COMMAND_STATEMENT_QUERY, query)
        _, inner = self._parse_any(buf)
        inner_fields = _parse_fields(inner)
        assert inner_fields[1].decode("utf-8") == query


class TestNetworkXIntegration:
    """Verify pandas → NetworkX path without a server."""

    def test_from_pandas_edgelist(self):
        import pandas as pd
        import networkx as nx

        edge_df = pd.DataFrame({"src": [0, 0, 1], "dst": [1, 2, 2], "weight": [1.0, 2.0, 3.0]})
        G = nx.from_pandas_edgelist(edge_df, source="src", target="dst",
                                     edge_attr=True, create_using=nx.DiGraph())
        assert G.number_of_nodes() == 3
        assert G.number_of_edges() == 3
        assert G[0][1]["weight"] == 1.0

    def test_set_node_attributes_from_df(self):
        import pandas as pd
        import networkx as nx

        G = nx.DiGraph()
        G.add_nodes_from([0, 1, 2])

        vertex_df = pd.DataFrame({"_gar_id": [0, 1, 2], "name": ["Alice", "Bob", "Carol"]})
        nx.set_node_attributes(G, vertex_df.set_index("_gar_id").to_dict(orient="index"))

        assert G.nodes[0]["name"] == "Alice"
        assert G.nodes[2]["name"] == "Carol"

    def test_full_graph_no_server(self):
        """Build a complete graph from DataFrames — the same path KnutClient uses."""
        import pandas as pd
        import networkx as nx

        vertex_df = pd.DataFrame({
            "_gar_id": [0, 1, 2],
            "name": ["Alice", "Bob", "Carol"],
            "age": [30, 25, 35],
        })
        edge_df = pd.DataFrame({
            "src": [0, 0, 1],
            "dst": [1, 2, 2],
            "since": ["2020", "2019", "2021"],
        })

        G = nx.from_pandas_edgelist(edge_df, source="src", target="dst",
                                     edge_attr=True, create_using=nx.DiGraph())
        nx.set_node_attributes(G, vertex_df.set_index("_gar_id").to_dict(orient="index"))

        assert G.nodes[0]["name"] == "Alice"
        assert G[0][1]["since"] == "2020"
        # Carol (node 2) has the highest in-degree — receives links from both 0 and 1.
        in_degrees = dict(G.in_degree())
        assert max(in_degrees, key=in_degrees.get) == 2