pyMC_Repeater/tests/test_tx_lock.py

"""
Tests for TX lock serialisation and duty-cycle TOCTOU fix.
Addresses the three concerns raised in PR 190 review:

  1. Concurrent delayed_sends must not interleave send_packet calls.
  2. Duty-cycle TOCTOU must be fixed: the second packet is dropped when the
     first consumes the airtime budget inside the lock.
  3. Local retry must NOT hold _tx_lock during the 1-second backoff sleep —
     other queued packets must be able to transmit during that window.

Run with:
    python -m pytest tests/test_tx_lock.py -v
or:
    python -m unittest tests.test_tx_lock
"""

import asyncio
import time
import unittest
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

# ---------------------------------------------------------------------------
# Minimal handler factory
# ---------------------------------------------------------------------------


def _make_handler():
    """
    Return a RepeaterHandler instance with all external I/O mocked.

    Uses __new__ + manual attribute injection to bypass StorageCollector,
    radio hardware, and other heavy dependencies that are irrelevant to the
    TX lock behaviour under test.
    """
    from repeater.engine import RepeaterHandler

    radio = MagicMock()
    radio.spreading_factor = 9
    radio.bandwidth = 62500
    radio.coding_rate = 5
    radio.preamble_length = 17
    radio.frequency = 915_000_000
    radio.tx_power = 14

    dispatcher = MagicMock()
    dispatcher.radio = radio
    dispatcher.local_identity = None
    dispatcher.send_packet = AsyncMock()

    h = RepeaterHandler.__new__(RepeaterHandler)
    h.config = {
        "repeater": {"mode": "forward", "cache_ttl": 3600, "send_advert_interval_hours": 0},
        "delays": {"tx_delay_factor": 1.0, "direct_tx_delay_factor": 0.5},
        "duty_cycle": {"enforcement_enabled": True, "max_airtime_per_minute": 3600},
        "storage": {},
        "mesh": {},
    }
    h.dispatcher = dispatcher
    h.airtime_mgr = MagicMock()
    h.airtime_mgr.can_transmit.return_value = (True, 0.0)
    h.airtime_mgr.calculate_airtime.return_value = 100.0
    h._tx_lock = asyncio.Lock()
    h.sent_flood_count = 0
    h.sent_direct_count = 0
    # Stub out _record_packet_sent so it doesn't touch packet.header constants
    h._record_packet_sent = MagicMock()
    return h


def _make_packet(size: int = 50) -> MagicMock:
    pkt = MagicMock()
    pkt.get_raw_length.return_value = size
    pkt.header = 0x00
    return pkt


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestTxLockSerialisation(unittest.IsolatedAsyncioTestCase):
    # ── Test 1: no interleaving ─────────────────────────────────────────────

    async def test_concurrent_sends_do_not_interleave(self):
        """
        Two delayed_sends with identical delays race to the radio.
        send_packet must never be called while another call is already
        in-flight — i.e. _tx_lock must gate them sequentially.
        """
        h = _make_handler()
        pkt = _make_packet()

        in_flight = [False]
        overlap_detected = [False]

        async def send_with_overlap_check(*args, **kwargs):
            if in_flight[0]:
                overlap_detected[0] = True
            in_flight[0] = True
            await asyncio.sleep(0.05)  # simulate ~50ms radio TX
            in_flight[0] = False

        h.dispatcher.send_packet.side_effect = send_with_overlap_check

        # Both tasks use the same tiny delay so their timers expire together.
        t1 = await h.schedule_retransmit(pkt, delay=0.01, airtime_ms=0)
        t2 = await h.schedule_retransmit(pkt, delay=0.01, airtime_ms=0)
        await asyncio.gather(t1, t2, return_exceptions=True)

        self.assertFalse(
            overlap_detected[0],
            "send_packet was entered while another call was already in-flight "
            "— _tx_lock is not serialising correctly",
        )
        self.assertEqual(
            h.dispatcher.send_packet.call_count, 2, "Expected exactly 2 send_packet calls"
        )

    # ── Test 2: TOCTOU duty-cycle fix ──────────────────────────────────────

    async def test_duty_cycle_toctou_is_fixed(self):
        """
        When two tasks both pass the advisory can_transmit() check in __call__
        before either has recorded airtime, the authoritative check inside
        _tx_lock must ensure only one of them actually transmits.

        Simulated here by making can_transmit return True for the first
        in-lock check and False for every subsequent one.
        """
        h = _make_handler()
        pkt = _make_packet()
        airtime_ms = 100.0

        # First lock-holder gets True; second gets False (budget consumed).
        allow = [True]

        def can_tx(ms):
            if allow[0]:
                allow[0] = False
                return (True, 0.0)
            return (False, 5.0)

        h.airtime_mgr.can_transmit.side_effect = can_tx

        # Both tasks start simultaneously (delay=0).
        t1 = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=airtime_ms)
        t2 = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=airtime_ms)
        await asyncio.gather(t1, t2, return_exceptions=True)

        self.assertEqual(
            h.dispatcher.send_packet.call_count,
            1,
            "Both packets were sent — duty-cycle TOCTOU race was NOT fixed",
        )

    # ── Test 3: retry backoff does not hold the lock ────────────────────────

    async def test_local_retry_releases_lock_during_backoff(self):
        """
        When a local_transmission send fails on the first attempt, the 1-second
        backoff sleep must happen with _tx_lock released.

        We schedule:
          - pkt_local: local_transmission=True, delay=0s, fails on attempt 1
          - pkt_other: local_transmission=False, delay=0.1s

        pkt_other fires at ~0.1s.  Without the fix, the backoff sleep holds
        the lock until ~1.0s, so pkt_other would have to wait.  With the fix,
        pkt_other sends freely at ~0.1s, well before pkt_local retries at ~1.0s.
        """
        h = _make_handler()
        pkt_local = _make_packet()
        pkt_other = _make_packet()

        send_times: dict[int, float] = {}
        first_local_call = [True]

        async def tracked_send(*args, **kwargs):
            pkt = args[0]
            if pkt is pkt_local and first_local_call[0]:
                first_local_call[0] = False
                raise RuntimeError("simulated transient radio failure")
            send_times[id(pkt)] = time.monotonic()

        h.dispatcher.send_packet.side_effect = tracked_send

        t_local = await h.schedule_retransmit(
            pkt_local, delay=0.0, airtime_ms=0, local_transmission=True
        )
        t_other = await h.schedule_retransmit(
            pkt_other, delay=0.1, airtime_ms=0, local_transmission=False
        )
        await asyncio.gather(t_local, t_other, return_exceptions=True)

        self.assertIn(id(pkt_other), send_times, "pkt_other was never sent")
        self.assertIn(id(pkt_local), send_times, "pkt_local retry was never sent")

        # pkt_other fires at ~0.1s; pkt_local retry fires at ~1.0s.
        # If the lock were held during backoff, pkt_other would block until ~1.0s
        # and would be recorded AFTER pkt_local's retry — the assertion below
        # would fail.
        self.assertLess(
            send_times[id(pkt_other)],
            send_times[id(pkt_local)],
            "pkt_other sent AFTER pkt_local retry — "
            "_tx_lock was still held during the 1-second backoff sleep",
        )

    # ── Test 4: non-local single-attempt re-raises on failure ──────────────

    async def test_non_local_failure_propagates(self):
        """A relayed (non-local) packet that fails send_packet raises immediately."""
        h = _make_handler()
        pkt = _make_packet()

        h.dispatcher.send_packet.side_effect = RuntimeError("radio error")

        task = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=0, local_transmission=False)
        with self.assertRaises(RuntimeError):
            await task

        # Only one attempt should have been made.
        self.assertEqual(h.dispatcher.send_packet.call_count, 1)

    # ── Test 5: duty-cycle check re-runs after backoff ──────────────────────

    async def test_duty_cycle_rechecked_on_retry(self):
        """
        If the duty cycle is exhausted during the 1-second backoff, the retry
        attempt must still be dropped — i.e. the duty-cycle gate runs on every
        lock acquisition, not just the first.
        """
        h = _make_handler()
        pkt = _make_packet()

        # First attempt: send_packet raises → triggers backoff.
        # Between attempts the budget is consumed, so the retry lock sees False.
        transmit_seq = iter([(True, 0.0), (False, 5.0)])
        h.airtime_mgr.can_transmit.side_effect = lambda ms: next(transmit_seq)

        send_calls = [0]

        async def failing_then_gone(*args, **kwargs):
            send_calls[0] += 1
            if send_calls[0] == 1:
                raise RuntimeError("transient failure")
            # Should not reach here on attempt 1 (gate rejects)
            pytest_fail = AssertionError("send_packet called on attempt 1 despite gate rejection")
            raise pytest_fail

        h.dispatcher.send_packet.side_effect = failing_then_gone

        task = await h.schedule_retransmit(
            pkt, delay=0.0, airtime_ms=100.0, local_transmission=True
        )
        result = await task
        self.assertFalse(result, "Duty-cycle drop should report TX failure")

        self.assertEqual(
            send_calls[0], 1, "send_packet called on retry despite duty-cycle rejection"
        )

    # ── Test 6: explicit bridge radio_id is honoured ──────────────────────

    async def test_preferred_tx_radio_id_passed_to_dispatcher(self):
        """When a packet carries _preferred_tx_radio_id, TX call uses radio_id."""
        h = _make_handler()
        pkt = _make_packet()

        task = await h.schedule_retransmit(
            pkt,
            delay=0.0,
            airtime_ms=0,
            local_transmission=False,
            preferred_tx_radio_id="link",
        )
        result = await task

        self.assertTrue(result)
        self.assertEqual(h.dispatcher.send_packet.call_count, 1)
        _, kwargs = h.dispatcher.send_packet.call_args
        self.assertEqual(kwargs.get("radio_id"), "link")
        self.assertEqual(kwargs.get("wait_for_ack"), False)


class TestLocalTxRadioLinkWait(unittest.IsolatedAsyncioTestCase):
    """A local TX waits out a radio-link outage instead of spending its retry on it."""

    async def test_waits_for_a_dropped_link_then_transmits(self):
        """A KISS modem that drops off USB re-handshakes in ~2s; hold the TX for it."""
        h = _make_handler()
        h.config["delays"]["local_tx_link_wait_seconds"] = 2.0
        pkt = _make_packet()

        h.dispatcher.radio.is_connected = False
        h.dispatcher.radio.is_degraded = True

        async def restore_link():
            await asyncio.sleep(0.25)
            h.dispatcher.radio.is_connected = True
            h.dispatcher.radio.is_degraded = False

        asyncio.create_task(restore_link())

        started = time.monotonic()
        task = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=0, local_transmission=True)
        result = await task
        elapsed = time.monotonic() - started

        self.assertTrue(result)
        self.assertEqual(h.dispatcher.send_packet.call_count, 1)  # no attempt wasted
        self.assertGreaterEqual(elapsed, 0.25)  # it really waited

    async def test_gives_up_when_the_link_never_returns(self):
        """A wedged link is bounded, not waited on forever, and never transmitted to."""
        h = _make_handler()
        h.config["delays"]["local_tx_link_wait_seconds"] = 0.3
        pkt = _make_packet()

        h.dispatcher.radio.is_connected = False
        h.dispatcher.radio.is_degraded = True

        started = time.monotonic()
        task = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=0, local_transmission=True)
        result = await task
        elapsed = time.monotonic() - started

        self.assertFalse(result)
        h.dispatcher.send_packet.assert_not_called()
        self.assertLess(elapsed, 2.0)

    async def test_radio_without_link_state_is_not_delayed(self):
        """Radios that track no link state (direct SPI) must keep their timing."""
        h = _make_handler()  # MagicMock radio: is_connected/is_degraded are mocks
        pkt = _make_packet()

        started = time.monotonic()
        task = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=0, local_transmission=True)
        result = await task
        elapsed = time.monotonic() - started

        self.assertTrue(result)
        self.assertLess(elapsed, 0.2)

    async def test_relayed_packet_does_not_wait(self):
        """Only local TX waits; a repeat held back that long is worse than dropped."""
        h = _make_handler()
        h.config["delays"]["local_tx_link_wait_seconds"] = 5.0
        pkt = _make_packet()

        h.dispatcher.radio.is_connected = False
        h.dispatcher.radio.is_degraded = True

        started = time.monotonic()
        task = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=0, local_transmission=False)
        await task
        elapsed = time.monotonic() - started

        h.dispatcher.send_packet.assert_called_once()
        self.assertLess(elapsed, 0.2)

    async def test_waits_for_the_selected_radio_not_the_default(self):
        """With a fabric, link state must follow preferred_tx_radio_id.

        Attribute access on a FabricRadio passes through to its default endpoint, so
        reading state off it would suppress a healthy selected radio (or wave through
        a dead one) whenever the default disagrees.
        """
        h = _make_handler()
        h.config["delays"]["local_tx_link_wait_seconds"] = 0.3
        pkt = _make_packet()

        up = SimpleNamespace(is_connected=True, is_degraded=False)
        down = SimpleNamespace(is_connected=False, is_degraded=True)
        endpoints = {"a": up, "b": down}
        fabric = MagicMock()
        fabric.radios = endpoints
        fabric.get_radio = lambda rid: endpoints[rid]
        h.dispatcher.radio.fabric = fabric
        # The default endpoint is down; the selected one is not.
        h.dispatcher.radio.is_connected = False
        h.dispatcher.radio.is_degraded = True

        started = time.monotonic()
        task = await h.schedule_retransmit(
            pkt, delay=0.0, airtime_ms=0, local_transmission=True, preferred_tx_radio_id="a"
        )
        result = await task
        elapsed = time.monotonic() - started

        self.assertTrue(result)
        h.dispatcher.send_packet.assert_called_once()
        self.assertLess(elapsed, 0.2)  # never waited on the default radio

    async def test_waits_when_the_selected_radio_is_the_down_one(self):
        """The converse: a down selected radio is waited for, default health aside."""
        h = _make_handler()
        h.config["delays"]["local_tx_link_wait_seconds"] = 0.3
        pkt = _make_packet()

        up = SimpleNamespace(is_connected=True, is_degraded=False)
        down = SimpleNamespace(is_connected=False, is_degraded=True)
        endpoints = {"a": up, "b": down}
        fabric = MagicMock()
        fabric.radios = endpoints
        fabric.get_radio = lambda rid: endpoints[rid]
        h.dispatcher.radio.fabric = fabric
        h.dispatcher.radio.is_connected = True  # default is healthy

        task = await h.schedule_retransmit(
            pkt, delay=0.0, airtime_ms=0, local_transmission=True, preferred_tx_radio_id="b"
        )
        result = await task

        self.assertFalse(result)
        h.dispatcher.send_packet.assert_not_called()

    async def test_multi_radio_fabric_without_a_target_does_not_wait(self):
        """The fabric picks at send time, so there is nothing to inspect: fail open."""
        h = _make_handler()
        h.config["delays"]["local_tx_link_wait_seconds"] = 5.0
        pkt = _make_packet()

        endpoints = {
            "a": SimpleNamespace(is_connected=False, is_degraded=True),
            "b": SimpleNamespace(is_connected=False, is_degraded=True),
        }
        fabric = MagicMock()
        fabric.radios = endpoints
        h.dispatcher.radio.fabric = fabric
        h.dispatcher.radio.is_connected = False

        started = time.monotonic()
        task = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=0, local_transmission=True)
        await task
        elapsed = time.monotonic() - started

        h.dispatcher.send_packet.assert_called_once()
        self.assertLess(elapsed, 0.2)

    async def test_retry_waits_when_the_link_drops_after_the_first_attempt(self):
        """The real sequence: send fails, the link goes, it returns, the retry lands."""
        h = _make_handler()
        h.config["delays"]["local_tx_link_wait_seconds"] = 2.0
        pkt = _make_packet()

        async def restore_link():
            await asyncio.sleep(0.25)
            h.dispatcher.radio.is_connected = True
            h.dispatcher.radio.is_degraded = False

        attempts = []

        async def send_side_effect(*args, **kwargs):
            attempts.append(1)
            if len(attempts) == 1:
                # The modem falls off the bus as the first attempt fails.
                h.dispatcher.radio.is_connected = False
                h.dispatcher.radio.is_degraded = True
                asyncio.create_task(restore_link())
                return False
            return True

        h.dispatcher.send_packet.side_effect = send_side_effect

        started = time.monotonic()
        task = await h.schedule_retransmit(pkt, delay=0.0, airtime_ms=0, local_transmission=True)
        result = await task
        elapsed = time.monotonic() - started

        self.assertTrue(result)
        self.assertEqual(len(attempts), 2)
        self.assertGreaterEqual(elapsed, 0.25)  # waited for the link
        self.assertLess(elapsed, 1.0)  # and skipped the flat 1s backoff

    def test_configured_wait_must_be_finite_and_bounded(self):
        """An infinite or absurd ceiling would hold the companion's command open."""
        h = _make_handler()
        for value, expected in (
            (float("inf"), h.DEFAULT_LOCAL_TX_LINK_WAIT_S),
            (float("nan"), h.DEFAULT_LOCAL_TX_LINK_WAIT_S),
            ("not a number", h.DEFAULT_LOCAL_TX_LINK_WAIT_S),
            (1e9, h.MAX_LOCAL_TX_LINK_WAIT_S),
            (-5, 0.0),
            (2.5, 2.5),
        ):
            h.config["delays"]["local_tx_link_wait_seconds"] = value
            self.assertEqual(h._local_tx_link_wait_s(), expected, msg=f"for {value!r}")


if __name__ == "__main__":
    unittest.main()