host.plugin_data now always holds rtt_ipv4/rtt_ipv6 history after the first heartbeat, which broke two pre-existing behaviors that assumed plugin_data reflects only real client-collected data: - The request_update ACK gate in handle_datagram() never fired again after a host's first heartbeat, so stale OS/agent-version info was never refreshed after a reconnect (only a client restart fixed it). - plugin_data.clear() on every UP transition wiped rtt_* history too, so a flaky host could never accumulate a useful RTT graph. Add _is_rtt_key()/_has_real_plugin_data() helpers so both spots treat rtt_* keys as synthetic: the gate now looks only at real plugin data, and the recovery clear only drops real plugin keys, preserving RTT history across reconnects. Also fixes two pre-existing E127 continuation-indent flake8 issues in tests/test_udp_rtt_history.py. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DEimzMv4Q5EjFg3hoiZ69T
124 lines
4.3 KiB
Python
124 lines
4.3 KiB
Python
"""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):
|
|
d = {"name": name, "interval": interval, "id": 0}
|
|
if rtt is not None:
|
|
d["rtt"] = rtt
|
|
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_recovery_clears_real_plugin_data_but_preserves_rtt_history():
|
|
"""Regression for Finding 2: host.plugin_data.clear() on recovery must
|
|
wipe real client-collected plugin data while leaving rtt_* history
|
|
samples intact.
|
|
"""
|
|
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
|
|
|
|
# Recovery heartbeat.
|
|
handle_datagram(_htb("rtt-hist-host5", rtt=13.0), ("127.0.0.1", 50000), transport, ctx)
|
|
|
|
assert "os_info" not in host.plugin_data
|
|
assert len(host.plugin_data["rtt_ipv4"]) == 4
|