heartbeat/tests/test_udp_presence.py

"""Host-level overdue alerting driven by client-reported protocol presence."""
import asyncio
import time

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


class _FakeTransport:
    def sendto(self, data, addr):
        pass


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


def _fresh(name):
    hbdclass.Host.hosts.pop(name, None)


def _send(name, addr, ctx, **fields):
    msg = parse_message(dicttos("HTB", {"name": name, "id": 0, **fields}))
    handle_datagram(msg, (addr, 50000), _FakeTransport(), ctx)


def _events(monkeypatch):
    events = []
    monkeypatch.setattr(udp, "eventlog", lambda u, level, msg, **kw: events.append((level, msg)))
    return events


def _dual_stack(name, ctx):
    _fresh(name)
    _send(name, "127.0.0.1", ctx)
    _send(name, "::1", ctx)
    return hbdclass.Host.hosts[name]


def test_flags_are_stored_on_host():
    ctx = _ctx()
    _fresh("pres-store")
    _send("pres-store", "127.0.0.1", ctx, p4=1, p6=0)
    h = hbdclass.Host.hosts["pres-store"]
    assert h.presence["IPv4"]["present"] is True
    assert h.presence["IPv6"]["present"] is False


def test_no_flags_leaves_presence_unknown():
    ctx = _ctx()
    _fresh("pres-legacy")
    _send("pres-legacy", "127.0.0.1", ctx)
    assert hbdclass.Host.hosts["pres-legacy"].presence == {}


def test_absent_flag_moves_overdue_family_to_absent_and_clears_alert(monkeypatch):
    _events(monkeypatch)
    ctx = _ctx()
    h = _dual_stack("pres-absent", ctx)
    v6 = h.connections["IPv6"]
    v6.state = C.OVERDUE
    udp._set_connectivity_alert(h, "IPv6", "CRITICAL")
    _send("pres-absent", "127.0.0.1", ctx, p4=1, p6=0)
    assert v6.state == C.ABSENT
    assert "connectivity.IPv6" not in h.alert_states
    assert h.status() == "up"


def test_absent_flag_does_not_touch_family_that_is_up():
    ctx = _ctx()
    h = _dual_stack("pres-up-stays", ctx)
    _send("pres-up-stays", "127.0.0.1", ctx, p4=1, p6=0)
    assert h.connections["IPv6"].state == C.UP          # traffic overrides the flag


def test_present_again_arms_grace_timer():
    ctx = _ctx()

    async def scenario():
        _fresh("pres-back")
        _send("pres-back", "127.0.0.1", ctx, interval=10)
        _send("pres-back", "::1", ctx, interval=10)
        h = hbdclass.Host.hosts["pres-back"]
        v6 = h.connections["IPv6"]
        v6.cancel_overdue_timer()
        v6.state = C.ABSENT
        _send("pres-back", "127.0.0.1", ctx, interval=10, p6=1)
        assert v6.state == C.UNKNOWN and v6.overdue_timer is not None
        for c in h.connections.values():
            c.cancel_overdue_timer()

    asyncio.run(scenario())


def test_one_family_overdue_is_simplex_all_overdue_is_critical(monkeypatch):
    events = _events(monkeypatch)
    ctx = _ctx()
    h = _dual_stack("pres-simplex", ctx)
    on_overdue, _ = udp._make_timer_callbacks("pres-simplex", h, ctx)

    asyncio.run(on_overdue(h.connections["IPv6"]))
    assert h.status() == "simplex"
    assert ("INFO", "IPv6 overdue (simplex)") in events
    assert not any(level == "CRITICAL" for level, _ in events)
    assert "connectivity.IPv6" not in h.alert_states

    asyncio.run(on_overdue(h.connections["IPv4"]))
    assert h.status() == "overdue"
    assert [e for e in events if e[0] == "CRITICAL"] == [("CRITICAL", "IPv4, IPv6 overdue")]
    assert {"connectivity.IPv4", "connectivity.IPv6"} <= set(h.alert_states)


def test_timer_expiry_on_absent_family_is_silent(monkeypatch):
    events = _events(monkeypatch)
    ctx = _ctx()
    h = _dual_stack("pres-quiet", ctx)
    _send("pres-quiet", "127.0.0.1", ctx, p4=1, p6=0)   # v6 still UP => left alone
    events.clear()
    on_overdue, _ = udp._make_timer_callbacks("pres-quiet", h, ctx)
    asyncio.run(on_overdue(h.connections["IPv6"]))
    assert h.connections["IPv6"].state == C.ABSENT
    assert events == []
    assert h.status() == "up"


def test_recovery_from_simplex_is_info_only(monkeypatch):
    events = _events(monkeypatch)
    ctx = _ctx()
    h = _dual_stack("pres-rec-simplex", ctx)
    v6 = h.connections["IPv6"]
    v6.state, v6.statetime = C.OVERDUE, time.time() - 100
    events.clear()
    _send("pres-rec-simplex", "::1", ctx)
    assert v6.state == C.UP
    assert [e[0] for e in events] == ["INFO"] and "back after" in events[0][1]


def test_recovery_from_overdue_is_recover_and_clears_all_alerts(monkeypatch):
    events = _events(monkeypatch)
    ctx = _ctx()
    h = _dual_stack("pres-rec-overdue", ctx)
    for afam in ("IPv4", "IPv6"):
        c = h.connections[afam]
        c.state, c.statetime = C.OVERDUE, time.time() - 100
        udp._set_connectivity_alert(h, afam, "CRITICAL")
    events.clear()
    _send("pres-rec-overdue", "127.0.0.1", ctx)
    assert [e[0] for e in events] == ["RECOVER"]
    assert h.status() == "simplex"
    assert "connectivity.IPv6" not in h.alert_states     # left overdue => all cleared
    assert "connectivity.IPv4" not in h.alert_states


def test_restore_skips_absent_and_sets_no_alert_for_simplex():
    ctx = _ctx()

    async def scenario():
        h = _dual_stack("pres-restore", ctx)
        h.connections["IPv6"].state = C.OVERDUE          # v4 UP + v6 OVERDUE => simplex
        udp.restore_connection_timers(hbdclass, ctx)
        assert "connectivity.IPv6" not in h.alert_states
        h.connections["IPv6"].cancel_overdue_timer()
        h.connections["IPv6"].state = C.ABSENT
        h.set_presence("IPv6", False, time.time())
        udp.restore_connection_timers(hbdclass, ctx)
        assert h.connections["IPv6"].state == C.ABSENT
        for c in h.connections.values():
            c.cancel_overdue_timer()

    asyncio.run(scenario())


def _past_drop(conn):
    conn.state, conn.statetime = C.OVERDUE, time.time() - (udp.DROPOVERDUE + 10)


def test_restore_past_drop_window_keeps_alert_single_family():
    ctx = _ctx()

    async def scenario():
        _fresh("pres-drop-single")
        _send("pres-drop-single", "127.0.0.1", ctx)
        h = hbdclass.Host.hosts["pres-drop-single"]
        _past_drop(h.connections["IPv4"])
        udp.restore_connection_timers(hbdclass, ctx)
        assert h.connections["IPv4"].state == C.UNKNOWN
        assert "connectivity.IPv4" in h.alert_states

    asyncio.run(scenario())


def test_restore_past_drop_window_keeps_alerts_dual_family():
    ctx = _ctx()

    async def scenario():
        h = _dual_stack("pres-drop-dual", ctx)
        for c in h.connections.values():
            _past_drop(c)
        udp.restore_connection_timers(hbdclass, ctx)
        assert {"connectivity.IPv4", "connectivity.IPv6"} <= set(h.alert_states)

    asyncio.run(scenario())


def test_v4_only_host_overdue_is_single_critical(monkeypatch):
    events = _events(monkeypatch)
    ctx = _ctx()
    _fresh("pres-v4only")
    _send("pres-v4only", "127.0.0.1", ctx)
    h = hbdclass.Host.hosts["pres-v4only"]
    on_overdue, _ = udp._make_timer_callbacks("pres-v4only", h, ctx)
    asyncio.run(on_overdue(h.connections["IPv4"]))
    assert [e for e in events if e[0] == "CRITICAL"] == [("CRITICAL", "IPv4 overdue")]
    assert "connectivity.IPv4" in h.alert_states
    asyncio.run(on_overdue(h.connections["IPv4"]))
    assert len([e for e in events if e[0] == "CRITICAL"]) == 1


def test_all_families_absent_v4_only_host_still_alerts_when_it_dies(monkeypatch):
    events = _events(monkeypatch)
    ctx = _ctx()
    _fresh("pres-allabsent")
    _send("pres-allabsent", "127.0.0.1", ctx, p4=0)
    h = hbdclass.Host.hosts["pres-allabsent"]
    on_overdue, _ = udp._make_timer_callbacks("pres-allabsent", h, ctx)
    asyncio.run(on_overdue(h.connections["IPv4"]))
    assert [e for e in events if e[0] == "CRITICAL"] == [("CRITICAL", "IPv4 overdue")]
    assert h.connections["IPv4"].state == C.OVERDUE
    assert h.status() == "overdue"
    assert "connectivity.IPv4" in h.alert_states


def test_last_present_family_is_not_parked_when_other_is_absent(monkeypatch):
    events = _events(monkeypatch)
    ctx = _ctx()
    h = _dual_stack("pres-lastfam", ctx)
    _send("pres-lastfam", "127.0.0.1", ctx, p4=0, p6=0)
    events.clear()
    on_overdue, _ = udp._make_timer_callbacks("pres-lastfam", h, ctx)
    asyncio.run(on_overdue(h.connections["IPv4"]))
    assert h.connections["IPv4"].state == C.ABSENT
    assert not any(level == "CRITICAL" for level, _ in events)
    assert h.status() == "up"
    asyncio.run(on_overdue(h.connections["IPv6"]))
    assert [e for e in events if e[0] == "CRITICAL"] == [("CRITICAL", "IPv6 overdue")]
    assert h.connections["IPv6"].state == C.OVERDUE
    assert h.status() == "overdue"