i2p-python / tests / test_data / test_message_router.py

"""Tests for I2NP message routing — InboundMessageHandler, OutboundMessageRouter, MessageDispatcher."""

import struct

import pytest


class TestInboundMessageHandlerRegistration:
    def test_register_and_dispatch_custom_handler(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()

        def custom(payload: bytes) -> bytes:
            return payload[::-1]

        handler.register(99, custom)
        result = handler.handle(99, b"abc")
        assert result == b"cba"

    def test_handle_returns_none_for_unregistered_type(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        result = handler.handle(255, b"anything")
        assert result is None


class TestInboundDatabaseStoreAndLookup:
    def test_store_and_lookup(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        key = b"\x01" * 32
        value = b"router info data"
        payload = key + value
        handler.on_database_store(payload)

        result = handler.on_database_lookup(key)
        assert result == value

    def test_lookup_missing_key_returns_none(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        key = b"\x02" * 32
        result = handler.on_database_lookup(key)
        assert result is None


class TestInboundDeliveryStatus:
    def test_extract_message_id(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        msg_id = 0xDEADBEEF
        payload = msg_id.to_bytes(4, "big") + b"\x00" * 8
        result = handler.on_delivery_status(payload)
        assert result == 0xDEADBEEF

    def test_extract_message_id_zero(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        payload = (0).to_bytes(4, "big") + b"\x00" * 8
        result = handler.on_delivery_status(payload)
        assert result == 0


class TestInboundTunnelData:
    def test_extract_tunnel_id_and_data(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        tunnel_id = 42
        data = b"tunnel payload here"
        payload = tunnel_id.to_bytes(4, "big") + data
        tid, tdata = handler.on_tunnel_data(payload)
        assert tid == 42
        assert tdata == data

    def test_extract_tunnel_id_empty_data(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        payload = (100).to_bytes(4, "big")
        tid, tdata = handler.on_tunnel_data(payload)
        assert tid == 100
        assert tdata == b""


class TestInboundGarlic:
    def test_on_garlic_returns_single_clove(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        payload = b"garlic clove data"
        cloves = handler.on_garlic(payload)
        assert cloves == [payload]


class TestInboundDefaultHandlers:
    def test_default_handlers_registered(self):
        from i2p_data.message_router import InboundMessageHandler

        handler = InboundMessageHandler()
        # DATABASE_STORE=1, DATABASE_LOOKUP=2, DELIVERY_STATUS=10,
        # GARLIC=11, TUNNEL_DATA=18, TUNNEL_GATEWAY=19
        assert 1 in handler._handlers
        assert 2 in handler._handlers
        assert 10 in handler._handlers
        assert 11 in handler._handlers
        assert 18 in handler._handlers
        assert 19 in handler._handlers


class TestOutboundMessageRouterLocal:
    def test_route_local(self):
        from i2p_data.message_router import OutboundMessageRouter

        router = OutboundMessageRouter()
        result = router.route(0, b"local payload")
        assert result == {"type": "local", "payload": b"local payload"}


class TestOutboundMessageRouterRouter:
    def test_route_router(self):
        from i2p_data.message_router import OutboundMessageRouter

        router = OutboundMessageRouter()
        rh = b"\xAA" * 32
        result = router.route(1, b"data", router_hash=rh)
        assert result == {"type": "router", "router_hash": rh, "payload": b"data"}


class TestOutboundMessageRouterTunnel:
    def test_route_tunnel(self):
        from i2p_data.message_router import OutboundMessageRouter

        router = OutboundMessageRouter()
        gw = b"\xBB" * 32
        result = router.route(2, b"tdata", tunnel_id=7, gateway=gw)
        assert result == {
            "type": "tunnel",
            "tunnel_id": 7,
            "gateway": gw,
            "payload": b"tdata",
        }


class TestOutboundMessageRouterDestination:
    def test_route_destination(self):
        from i2p_data.message_router import OutboundMessageRouter

        router = OutboundMessageRouter()
        dest = b"\xCC" * 32
        result = router.route(3, b"dest_data", destination=dest)
        assert result == {
            "type": "destination",
            "destination": dest,
            "payload": b"dest_data",
        }


class TestOutboundMessageRouterUnknown:
    def test_route_unknown_raises_value_error(self):
        from i2p_data.message_router import OutboundMessageRouter

        router = OutboundMessageRouter()
        with pytest.raises(ValueError, match="Unknown delivery type"):
            router.route(99, b"payload")


class TestMessageDispatcher:
    def test_dispatch_inbound(self):
        from i2p_data.message_router import (
            InboundMessageHandler,
            MessageDispatcher,
            OutboundMessageRouter,
        )

        inbound = InboundMessageHandler()
        outbound = OutboundMessageRouter()
        dispatcher = MessageDispatcher(inbound, outbound)

        # Dispatch delivery status (type 10)
        msg_id = 0x12345678
        payload = msg_id.to_bytes(4, "big") + b"\x00" * 8
        result = dispatcher.dispatch_inbound(10, payload)
        assert result == 0x12345678

    def test_dispatch_outbound(self):
        from i2p_data.message_router import (
            InboundMessageHandler,
            MessageDispatcher,
            OutboundMessageRouter,
        )

        inbound = InboundMessageHandler()
        outbound = OutboundMessageRouter()
        dispatcher = MessageDispatcher(inbound, outbound)

        result = dispatcher.dispatch_outbound(0, b"payload")
        assert result == {"type": "local", "payload": b"payload"}

    def test_register_inbound_handler(self):
        from i2p_data.message_router import (
            InboundMessageHandler,
            MessageDispatcher,
            OutboundMessageRouter,
        )

        inbound = InboundMessageHandler()
        outbound = OutboundMessageRouter()
        dispatcher = MessageDispatcher(inbound, outbound)

        def my_handler(payload: bytes) -> str:
            return "handled"

        dispatcher.register_inbound_handler(200, my_handler)
        result = dispatcher.dispatch_inbound(200, b"test")
        assert result == "handled"