heartbeat/tests/test_udp_rtt_history.py

"""Tests for RTT history capture in udp.py's handle_datagram."""
import time

from hbd.common.proto import dicttos
from hbd.server import hbdclass
from hbd.server.udp import handle_datagram, parse_message


class _FakeTransport:
    def __init__(self):
        self.sent = []

    def sendto(self, data, addr):
        self.sent.append((data, addr))


def _htb(name, rtt=None, interval=0, boot=0):
    d = {"name": name, "interval": interval, "id": 0}
    if rtt is not None:
        d["rtt"] = rtt
    if boot:
        d["boot"] = boot
    return parse_message(dicttos("HTB", d))


def _base_ctx():
    return {
        "config": {},
        "hbdclass": hbdclass,
        "msg_to_websockets": None,
        "DEBUG": 0,
        "verbose": False,
    }


def test_handle_datagram_records_rtt_history_for_new_connection():
    hbdclass.Host.hosts.pop("rtt-hist-host", None)
    handle_datagram(_htb("rtt-hist-host", rtt=42.5), ("127.0.0.1", 50000),
                    _FakeTransport(), _base_ctx())

    host = hbdclass.Host.hosts["rtt-hist-host"]
    samples = host.plugin_data.get("rtt_ipv4")
    assert samples is not None
    assert len(samples) == 1
    ts, data = samples[0]
    assert data == {"rtt": 42.5}
    assert isinstance(ts, float)


def test_handle_datagram_appends_rtt_history_across_heartbeats():
    hbdclass.Host.hosts.pop("rtt-hist-host2", None)
    transport = _FakeTransport()
    ctx = _base_ctx()
    handle_datagram(_htb("rtt-hist-host2", rtt=10.0), ("127.0.0.1", 50000), transport, ctx)
    handle_datagram(_htb("rtt-hist-host2", rtt=20.0), ("127.0.0.1", 50000), transport, ctx)

    host = hbdclass.Host.hosts["rtt-hist-host2"]
    samples = host.plugin_data["rtt_ipv4"]
    assert [d["rtt"] for _, d in samples] == [10.0, 20.0]


def test_handle_datagram_skips_rtt_history_when_rtt_missing():
    hbdclass.Host.hosts.pop("rtt-hist-host3", None)
    handle_datagram(_htb("rtt-hist-host3", rtt=None), ("127.0.0.1", 50000),
                    _FakeTransport(), _base_ctx())

    host = hbdclass.Host.hosts["rtt-hist-host3"]
    assert "rtt_ipv4" not in host.plugin_data


def test_request_update_fires_on_recovery_even_with_rtt_history():
    """Regression for Finding 1: rtt_* keys must not permanently disable the
    request_update gate. A connection recovering from a non-UP state must
    still be asked to resend real plugin data, even though rtt_ipv4 already
    holds samples from before the drop.
    """
    hbdclass.Host.hosts.pop("rtt-hist-host4", None)
    transport = _FakeTransport()
    ctx = _base_ctx()

    # First heartbeat: brand-new host, no plugin data at all yet.
    handle_datagram(_htb("rtt-hist-host4", rtt=15.0), ("127.0.0.1", 50000), transport, ctx)
    host = hbdclass.Host.hosts["rtt-hist-host4"]
    assert host.plugin_data.get("rtt_ipv4")  # rtt history now non-empty

    # Simulate a recovery: connection was dropped (e.g. OVERDUE->UP after a
    # missed heartbeat) and is about to come back UP on the next heartbeat.
    conn = host.connections["IPv4"]
    conn.state = hbdclass.Connection.DOWN

    transport.sent.clear()
    handle_datagram(_htb("rtt-hist-host4", rtt=16.0), ("127.0.0.1", 50000), transport, ctx)

    ack_data, _ = transport.sent[0]
    ack = parse_message(ack_data)
    assert ack.get("request_update")


def test_ordinary_recovery_preserves_real_plugin_data_and_rtt_history():
    """An ordinary reconnect (no boot flag) — e.g. OVERDUE->UP after a
    transient network blip — must NOT wipe already-collected real plugin
    data (e.g. cpu_monitor, os_info) or rtt_* history. Only an actual
    client reboot invalidates that data (see the boot-flag test below).
    """
    hbdclass.Host.hosts.pop("rtt-hist-host5", None)
    transport = _FakeTransport()
    ctx = _base_ctx()

    for rtt in (10.0, 11.0, 12.0):
        handle_datagram(_htb("rtt-hist-host5", rtt=rtt), ("127.0.0.1", 50000), transport, ctx)

    host = hbdclass.Host.hosts["rtt-hist-host5"]
    assert len(host.plugin_data["rtt_ipv4"]) == 3

    # Simulate a drop, and pretend the client had previously sent real
    # plugin data (collected before the connection went down).
    conn = host.connections["IPv4"]
    conn.state = hbdclass.Connection.DOWN
    host.add_plugin_data("os_info", {"os": "linux"}, timestamp=time.time())
    assert "os_info" in host.plugin_data

    # Ordinary recovery heartbeat — no boot flag.
    handle_datagram(_htb("rtt-hist-host5", rtt=13.0), ("127.0.0.1", 50000), transport, ctx)

    assert "os_info" in host.plugin_data
    assert len(host.plugin_data["rtt_ipv4"]) == 4


def test_boot_recovery_clears_real_plugin_data_but_preserves_rtt_history():
    """A recovery heartbeat carrying the boot flag (client process actually
    restarted) must still wipe stale real plugin data, while rtt_* history
    (still a valid measurement, unaffected by a client reboot) survives.
    """
    hbdclass.Host.hosts.pop("rtt-hist-host6", None)
    transport = _FakeTransport()
    ctx = _base_ctx()

    for rtt in (10.0, 11.0, 12.0):
        handle_datagram(_htb("rtt-hist-host6", rtt=rtt), ("127.0.0.1", 50000), transport, ctx)

    host = hbdclass.Host.hosts["rtt-hist-host6"]
    assert len(host.plugin_data["rtt_ipv4"]) == 3

    conn = host.connections["IPv4"]
    conn.state = hbdclass.Connection.DOWN
    host.add_plugin_data("os_info", {"os": "linux"}, timestamp=time.time())
    assert "os_info" in host.plugin_data

    # Recovery heartbeat with boot=1 — client process actually restarted.
    handle_datagram(_htb("rtt-hist-host6", rtt=13.0, boot=1), ("127.0.0.1", 50000), transport, ctx)

    assert "os_info" not in host.plugin_data
    assert len(host.plugin_data["rtt_ipv4"]) == 4