test_anyscan_rate_controller.py

"""Unit + smoke tests for the AIMD port-scan rate controller.

Run via ``python3 -m unittest test_anyscan_rate_controller -v`` from the
anyscan repo root. Pure-Python; no scanner binary or NIC required.
"""

from __future__ import annotations

import io
import json
import os
import tempfile
import unittest
from dataclasses import dataclass
from pathlib import Path

import anyscan_rate_controller as rc


def make_measurement(
    *,
    set_rate: int,
    achieved_pps: float,
    elapsed_seconds: float = 30.0,
    tx_dropped_delta: int = 0,
    heartbeat_max_latency_ms: int = 0,
    scanner_finished_naturally: bool = False,
) -> rc.WindowMeasurement:
    return rc.WindowMeasurement(
        set_rate=set_rate,
        elapsed_seconds=elapsed_seconds,
        tx_packets_delta=int(achieved_pps * elapsed_seconds),
        tx_dropped_delta=tx_dropped_delta,
        heartbeat_max_latency_ms=heartbeat_max_latency_ms,
        scanner_finished_naturally=scanner_finished_naturally,
    )


class AimdMathTests(unittest.TestCase):
    def setUp(self) -> None:
        self.policy = rc.AimdPolicy()

    def test_clean_window_bumps_additively(self) -> None:
        measurement = make_measurement(set_rate=500_000, achieved_pps=480_000)
        self.assertEqual(rc.classify_window(measurement, self.policy), rc.CLEAN)
        next_rate = rc.compute_next_rate(self.policy, 500_000, measurement)
        self.assertEqual(next_rate, 700_000)

    def test_slip_due_to_drops_halves(self) -> None:
        measurement = make_measurement(
            set_rate=2_000_000,
            achieved_pps=1_900_000,
            tx_dropped_delta=42,
        )
        self.assertEqual(rc.classify_window(measurement, self.policy), rc.SLIP)
        next_rate = rc.compute_next_rate(self.policy, 2_000_000, measurement)
        self.assertEqual(next_rate, 1_000_000)

    def test_slip_due_to_heartbeat_halves(self) -> None:
        measurement = make_measurement(
            set_rate=3_000_000,
            achieved_pps=2_800_000,
            heartbeat_max_latency_ms=6_000,
        )
        self.assertEqual(rc.classify_window(measurement, self.policy), rc.SLIP)
        next_rate = rc.compute_next_rate(self.policy, 3_000_000, measurement)
        self.assertEqual(next_rate, 1_500_000)

    def test_slip_due_to_starved_achieved_rate(self) -> None:
        # set 5M but achieved only 1.5M — rate-limit overhead tax that PR #55
        # observed. Should be classified as slip even with no drops.
        measurement = make_measurement(set_rate=5_000_000, achieved_pps=1_500_000)
        self.assertEqual(rc.classify_window(measurement, self.policy), rc.SLIP)
        next_rate = rc.compute_next_rate(self.policy, 5_000_000, measurement)
        self.assertEqual(next_rate, 2_500_000)

    def test_natural_finish_skips_achieved_floor(self) -> None:
        # Tiny target ranges finish in milliseconds — achieved/set will look
        # tiny but that just means we ran out of work, not that we're being
        # throttled. Don't punish that.
        measurement = make_measurement(
            set_rate=2_000_000,
            achieved_pps=80_000,
            elapsed_seconds=0.5,
            scanner_finished_naturally=True,
        )
        self.assertEqual(rc.classify_window(measurement, self.policy), rc.CLEAN)

    def test_ceiling_clamp_on_clean(self) -> None:
        policy = rc.AimdPolicy(ceiling=1_000_000)
        measurement = make_measurement(set_rate=950_000, achieved_pps=940_000)
        self.assertEqual(rc.compute_next_rate(policy, 950_000, measurement), 1_000_000)

    def test_floor_clamp_on_slip(self) -> None:
        policy = rc.AimdPolicy(floor=300_000, multiplicative_factor=0.25)
        measurement = make_measurement(
            set_rate=400_000,
            achieved_pps=10_000,
            tx_dropped_delta=10,
        )
        self.assertEqual(rc.compute_next_rate(policy, 400_000, measurement), 300_000)

    def test_clamp_rate_pulls_in_outside_band(self) -> None:
        policy = rc.AimdPolicy(floor=100_000, ceiling=4_000_000)
        self.assertEqual(rc.clamp_rate(50_000, policy), 100_000)
        self.assertEqual(rc.clamp_rate(10_000_000, policy), 4_000_000)
        self.assertEqual(rc.clamp_rate(2_000_000, policy), 2_000_000)

    def test_invalid_policy_rejected(self) -> None:
        with self.assertRaises(ValueError):
            rc.AimdPolicy(floor=0)
        with self.assertRaises(ValueError):
            rc.AimdPolicy(floor=500_000, ceiling=400_000)
        with self.assertRaises(ValueError):
            rc.AimdPolicy(multiplicative_factor=0.0)
        with self.assertRaises(ValueError):
            rc.AimdPolicy(multiplicative_factor=1.0)


class CalibrationStoreTests(unittest.TestCase):
    def test_load_returns_empty_when_missing(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            store = rc.RateCalibrationStore(Path(tmpdir) / "rate-calibration.json")
            self.assertEqual(store.load(), {})
            self.assertIsNone(store.lookup("eth0"))

    def test_store_round_trip(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "rate-calibration.json"
            store = rc.RateCalibrationStore(path)
            store.store("eth0", 1_700_000, now_iso="2026-04-27T00:00:00Z")
            store.store("eth1", 800_000, now_iso="2026-04-27T00:01:00Z")
            self.assertEqual(store.lookup("eth0").learned_rate, 1_700_000)
            self.assertEqual(store.lookup("eth1").learned_rate, 800_000)
            payload = json.loads(path.read_text())
            self.assertEqual(payload["version"], 1)
            self.assertIn("eth0", payload["interfaces"])

    def test_corrupt_file_treated_as_empty(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "rate-calibration.json"
            path.write_text("not json{}")
            store = rc.RateCalibrationStore(path)
            self.assertEqual(store.load(), {})

    def test_zero_or_negative_rates_skipped(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "rate-calibration.json"
            store = rc.RateCalibrationStore(path)
            store.store("eth0", 0)
            store.store("eth0", -100)
            self.assertFalse(path.exists())

    def test_concurrent_writes_from_multiple_processes_all_persist(self) -> None:
        """Multi-NIC parent spawns N shard children; each child's controller
        writes to the SAME calibration file when it converges. Pre-fix the
        children clobbered each other (shared tmp filename + read-modify-write
        race) so only one shard's calibration survived. Post-fix all shards'
        entries land.
        """

        import multiprocessing as mp

        def _store_in_subprocess(path_str: str, interface: str, rate: int) -> None:
            # Re-import inside the worker so the subprocess has a clean
            # module state (no shared file handles across fork).
            from pathlib import Path as _Path
            import anyscan_rate_controller as _rc

            store = _rc.RateCalibrationStore(_Path(path_str))
            store.store(interface, rate, now_iso="2026-04-27T12:00:00Z")

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "rate-calibration.json"
            # Six shards, fired simultaneously. Each writes a distinct
            # interface; the per-interface rate value is the rate the
            # synthetic AIMD loop "converged" to.
            interfaces = [(f"eth{i}", 1_000_000 + i * 100_000) for i in range(6)]
            ctx = mp.get_context("fork")
            workers = [
                ctx.Process(
                    target=_store_in_subprocess,
                    args=(str(path), iface, rate),
                )
                for iface, rate in interfaces
            ]
            for w in workers:
                w.start()
            for w in workers:
                w.join(timeout=10)
                self.assertEqual(w.exitcode, 0, msg=f"worker {w.pid} failed")

            store = rc.RateCalibrationStore(path)
            entries = store.load()
            # Pre-fix: only one entry survived (last-writer-wins on the
            # shared tmp + clobbered .json). Post-fix: all six landed.
            self.assertEqual(
                set(entries.keys()),
                {iface for iface, _ in interfaces},
                msg=f"expected all 6 interfaces persisted, got {list(entries.keys())}",
            )
            for iface, rate in interfaces:
                self.assertEqual(entries[iface].learned_rate, rate)

    def test_concurrent_writes_leave_no_dangling_tmp_files(self) -> None:
        """The pre-fix race left orphan .tmp files on disk when the second
        os.replace failed (source already moved). Post-fix every writer
        cleans up its own per-pid tmp regardless of outcome.
        """

        import multiprocessing as mp

        def _store_in_subprocess(path_str: str, interface: str, rate: int) -> None:
            from pathlib import Path as _Path
            import anyscan_rate_controller as _rc

            store = _rc.RateCalibrationStore(_Path(path_str))
            store.store(interface, rate)

        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "rate-calibration.json"
            ctx = mp.get_context("fork")
            workers = [
                ctx.Process(
                    target=_store_in_subprocess,
                    args=(str(path), f"eth{i}", 500_000 + i * 100_000),
                )
                for i in range(8)
            ]
            for w in workers:
                w.start()
            for w in workers:
                w.join(timeout=10)
                self.assertEqual(w.exitcode, 0)

            # No stray .tmp.* files: success and failure paths both clean
            # up after themselves.
            stray = list(Path(tmpdir).glob("rate-calibration.json.tmp*"))
            self.assertEqual(stray, [], msg=f"orphan tmp files left: {stray}")
            # The final file is valid JSON we can round-trip.
            store = rc.RateCalibrationStore(path)
            self.assertEqual(len(store.load()), 8)


@dataclass
class StubWindow:
    achieved_pps_for_set: dict[int, float]
    drops_for_set: dict[int, int]
    heartbeat_for_set: dict[int, int]
    natural_finish_after_window: int = 100  # effectively never


class StubRunner(rc.WindowRunner):
    def __init__(self, scenario: StubWindow, *, profile: callable | None = None) -> None:
        self._scenario = scenario
        self._profile = profile
        self._invocation = 0

    def run(
        self,
        *,
        rate: int,
        window_seconds: float,
        is_first_window: bool,
    ) -> rc.WindowMeasurement:
        self._invocation += 1
        if self._profile is not None:
            achieved, dropped, heartbeat = self._profile(rate)
        else:
            achieved = self._scenario.achieved_pps_for_set.get(rate, 0)
            dropped = self._scenario.drops_for_set.get(rate, 0)
            heartbeat = self._scenario.heartbeat_for_set.get(rate, 0)
        finished = self._invocation >= self._scenario.natural_finish_after_window
        return rc.WindowMeasurement(
            set_rate=rate,
            elapsed_seconds=window_seconds,
            tx_packets_delta=int(achieved * window_seconds),
            tx_dropped_delta=dropped,
            heartbeat_max_latency_ms=heartbeat,
            scanner_finished_naturally=finished,
        )


class ConvergenceSmokeTests(unittest.TestCase):
    """Synthetic convergence test using PR #55's measured numbers as ground truth."""

    def _profile_from_pr55(self, rate: int) -> tuple[float, int, int]:
        # rate=100k achieved 175k (cap cosmetic below natural)
        # rate=500k achieved ~480k (estimated linear region)
        # rate=700k achieved ~680k
        # rate=900k achieved ~880k
        # rate=1M  achieved 1_760_000 (sweet spot kicked in here in real bench)
        # rate=2M  achieved 2_220_000 (max)
        # rate=2.2M+ rate-limit overhead starts hurting; achieved drops
        # rate=5M  achieved 1_590_000 (rate-limit overhead hurts)
        natural_ceiling = 2_220_000
        sweet_spot = 1_700_000
        if rate <= 100_000:
            achieved = float(rate) * 1.75
            return achieved, 0, 50
        if rate <= 1_000_000:
            achieved = min(float(rate) * 0.97, natural_ceiling)
            return achieved, 0, 100
        if rate <= 2_000_000:
            achieved = min(natural_ceiling - (rate - 2_000_000) * 0.0, sweet_spot + (rate - 1_000_000) * 0.5)
            return achieved, 0, 800
        if rate <= 3_000_000:
            # Some packets dropped as we cross the host's natural ceiling.
            achieved = natural_ceiling - (rate - 2_000_000) * 0.4
            return achieved, int((rate - 2_000_000) * 0.05), 4_500
        # Above 3M: rate-limit overhead AND drops AND scheduler starvation.
        return 1_500_000.0, int((rate - 3_000_000) * 0.1), 7_000

    def test_converges_into_sweet_spot_within_a_handful_of_windows(self) -> None:
        policy = rc.AimdPolicy(window_seconds=30)
        scenario = StubWindow(
            achieved_pps_for_set={},
            drops_for_set={},
            heartbeat_for_set={},
        )
        runner = StubRunner(scenario, profile=self._profile_from_pr55)
        sink = io.StringIO()
        controller = rc.RateController(
            options=rc.ControllerOptions(
                policy=policy,
                window_seconds=float(policy.window_seconds),
                interface="eth0",
                starting_rate=500_000,
                calibration=None,
                persist_on_clean=False,
            ),
            runner=runner,
            log_sink=sink,
            max_windows=10,
        )
        reports = controller.run()

        self.assertEqual(len(reports), 10)
        first_set_rates = [r.set_rate for r in reports[:4]]
        self.assertEqual(first_set_rates, [500_000, 700_000, 900_000, 1_100_000])

        # Within the cap of 10 windows we should have already touched a rate
        # that exercises the sweet-spot region (>= 1.5M) and never explored
        # the dangerous-overhead zone (>= 4M).
        self.assertTrue(any(r.set_rate >= 1_500_000 for r in reports))
        self.assertTrue(all(r.set_rate <= 4_000_000 for r in reports))

        # Settled rate (last window) should be in the productive range —
        # not stuck at the floor and not parked at the ceiling.
        settled = reports[-1].set_rate
        self.assertGreaterEqual(settled, 900_000)
        self.assertLessEqual(settled, 4_000_000)

    def test_clean_only_scenario_walks_up_to_ceiling(self) -> None:
        policy = rc.AimdPolicy(
            floor=100_000,
            ceiling=1_000_000,
            additive_step=200_000,
            window_seconds=30,
        )

        def profile(rate: int) -> tuple[float, int, int]:
            return float(rate) * 0.95, 0, 50

        scenario = StubWindow(
            achieved_pps_for_set={},
            drops_for_set={},
            heartbeat_for_set={},
        )
        runner = StubRunner(scenario, profile=profile)
        controller = rc.RateController(
            options=rc.ControllerOptions(
                policy=policy,
                window_seconds=float(policy.window_seconds),
                interface=None,
                starting_rate=200_000,
                calibration=None,
                persist_on_clean=False,
            ),
            runner=runner,
            log_sink=io.StringIO(),
            max_windows=8,
        )
        reports = controller.run()
        rates = [r.set_rate for r in reports]
        self.assertEqual(
            rates,
            [200_000, 400_000, 600_000, 800_000, 1_000_000, 1_000_000, 1_000_000, 1_000_000],
        )

    def test_persists_max_clean_rate_after_run(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            calibration = rc.RateCalibrationStore(
                Path(tmpdir) / "rate-calibration.json"
            )
            policy = rc.AimdPolicy(window_seconds=30)
            runner = StubRunner(
                StubWindow(
                    achieved_pps_for_set={},
                    drops_for_set={},
                    heartbeat_for_set={},
                ),
                profile=self._profile_from_pr55,
            )
            controller = rc.RateController(
                options=rc.ControllerOptions(
                    policy=policy,
                    window_seconds=float(policy.window_seconds),
                    interface="eth0",
                    starting_rate=500_000,
                    calibration=calibration,
                ),
                runner=runner,
                log_sink=io.StringIO(),
                max_windows=8,
            )
            controller.run()
            entry = calibration.lookup("eth0")
            self.assertIsNotNone(entry)
            self.assertGreaterEqual(entry.learned_rate, 900_000)


class ScannerFailureTests(unittest.TestCase):
    def test_mid_window_failure_aborts_loop(self) -> None:
        class CrashingRunner(rc.WindowRunner):
            def __init__(self) -> None:
                self.calls = 0

            def run(
                self, *, rate: int, window_seconds: float, is_first_window: bool
            ) -> rc.WindowMeasurement:
                self.calls += 1
                return rc.WindowMeasurement(
                    set_rate=rate,
                    elapsed_seconds=1.0,
                    tx_packets_delta=0,
                    tx_dropped_delta=0,
                    heartbeat_max_latency_ms=0,
                    scanner_finished_naturally=False,
                    scanner_exit_code=139,
                )

        runner = CrashingRunner()
        controller = rc.RateController(
            options=rc.ControllerOptions(
                policy=rc.AimdPolicy(window_seconds=30),
                window_seconds=30.0,
                interface=None,
                starting_rate=500_000,
                calibration=None,
            ),
            runner=runner,
            log_sink=io.StringIO(),
            max_windows=5,
        )
        with self.assertRaises(rc.ScannerWindowError) as ctx:
            controller.run()
        self.assertEqual(ctx.exception.exit_code, 139)
        self.assertEqual(runner.calls, 1)


class JitterMonitorTests(unittest.TestCase):
    def test_reset_clears_recorded_gap(self) -> None:
        monitor = rc.JitterMonitor(interval_seconds=0.01)
        monitor.start()
        try:
            import time as _t

            _t.sleep(0.1)
            self.assertGreaterEqual(monitor.max_gap_ms(), 0)
            monitor.reset()
            self.assertEqual(monitor.max_gap_ms(), 0)
        finally:
            monitor.stop()


class NicStatsReaderTests(unittest.TestCase):
    def test_reads_counters_from_synthetic_root(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            iface_root = Path(tmpdir) / "eth0" / "statistics"
            iface_root.mkdir(parents=True)
            (iface_root / "tx_packets").write_text("123\n")
            (iface_root / "tx_dropped").write_text("4\n")
            reader = rc.NicStatsReader("eth0", root=tmpdir)
            counters = reader.read()
            self.assertEqual(counters.tx_packets, 123)
            self.assertEqual(counters.tx_dropped, 4)

    def test_missing_files_return_zero(self) -> None:
        reader = rc.NicStatsReader("ghost0", root="/nonexistent")
        counters = reader.read()
        self.assertEqual(counters.tx_packets, 0)
        self.assertEqual(counters.tx_dropped, 0)


class EnvHelpersTests(unittest.TestCase):
    def test_policy_from_env_round_trip(self) -> None:
        env = {
            "ANYSCAN_RATE_FLOOR": "150000",
            "ANYSCAN_RATE_CEILING": "3500000",
            "ANYSCAN_RATE_ADDITIVE_STEP": "250000",
            "ANYSCAN_RATE_MULTIPLICATIVE_FACTOR": "0.25",
            "ANYSCAN_RATE_WINDOW_SECONDS": "45",
            "ANYSCAN_HEARTBEAT_LATENCY_THRESHOLD_MS": "3000",
        }
        policy = rc.policy_from_env(env)
        self.assertEqual(policy.floor, 150_000)
        self.assertEqual(policy.ceiling, 3_500_000)
        self.assertEqual(policy.additive_step, 250_000)
        self.assertAlmostEqual(policy.multiplicative_factor, 0.25)
        self.assertEqual(policy.window_seconds, 45)
        self.assertEqual(policy.heartbeat_latency_threshold_ms, 3_000)

    def test_env_flag_default_when_missing(self) -> None:
        self.assertTrue(rc.env_flag({}, "MISSING", default=True))
        self.assertFalse(rc.env_flag({"X": "no"}, "X", default=True))
        self.assertTrue(rc.env_flag({"X": "1"}, "X", default=False))

    def test_resolve_starting_rate_prefers_calibration(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            calibration = rc.RateCalibrationStore(Path(tmpdir) / "calib.json")
            calibration.store("eth0", 1_700_000)
            policy = rc.AimdPolicy()
            resolved = rc.resolve_starting_rate(
                policy=policy,
                interface="eth0",
                calibration=calibration,
                fallback_rate=500_000,
            )
            self.assertEqual(resolved, 1_700_000)

    def test_resolve_starting_rate_clamps_calibration_into_band(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            calibration = rc.RateCalibrationStore(Path(tmpdir) / "calib.json")
            calibration.store("eth0", 99_000_000)
            policy = rc.AimdPolicy(ceiling=4_000_000)
            resolved = rc.resolve_starting_rate(
                policy=policy,
                interface="eth0",
                calibration=calibration,
                fallback_rate=500_000,
            )
            self.assertEqual(resolved, 4_000_000)


class CpuVsNetworkSlipTests(unittest.TestCase):
    """Verify CPU-pressure scenarios stop cratering the rate (improvement #1)."""

    def setUp(self) -> None:
        self.policy = rc.AimdPolicy()

    def _load(self, *, load: float, vcpu: int = 8) -> rc.SystemLoad:
        return rc.SystemLoad(load_average_1min=load, vcpu_count=vcpu)

    def test_cpu_pressure_holds_rate_instead_of_halving(self) -> None:
        # 8 vCPUs, load 8.5 (load/vcpu=1.06 > 0.8) AND heartbeat slip,
        # tx_dropped=0, achieved == set: pure CPU starvation.
        measurement = make_measurement(
            set_rate=4_000_000,
            achieved_pps=3_900_000,
            tx_dropped_delta=0,
            heartbeat_max_latency_ms=6_000,
        )
        load = self._load(load=8.5, vcpu=8)
        self.assertEqual(
            rc.classify_window(measurement, self.policy, system_load=load),
            rc.SLIP_CPU,
        )
        self.assertEqual(
            rc.compute_next_rate(self.policy, 4_000_000, measurement, system_load=load),
            4_000_000,  # held, not halved
        )

    def test_network_pressure_still_halves_rate_with_low_load(self) -> None:
        # tx_dropped > 0, load low (no CPU pressure): pure network slip.
        measurement = make_measurement(
            set_rate=4_000_000,
            achieved_pps=3_900_000,
            tx_dropped_delta=10_000,
            heartbeat_max_latency_ms=200,
        )
        load = self._load(load=2.0, vcpu=8)  # 0.25 per vcpu, well under threshold
        self.assertEqual(
            rc.classify_window(measurement, self.policy, system_load=load),
            rc.SLIP_NETWORK,
        )
        self.assertEqual(
            rc.compute_next_rate(self.policy, 4_000_000, measurement, system_load=load),
            2_000_000,  # halved
        )

    def test_both_pressure_picks_dominant_via_drop_ratio(self) -> None:
        # Significant drop ratio: NIC genuinely overrun, network is dominant
        # cause even though CPU is also pegged. Halve.
        big_drops = make_measurement(
            set_rate=4_000_000,
            achieved_pps=3_900_000,
            tx_dropped_delta=200_000,  # 200k of ~117M packets — 0.17% > 0.1% threshold
            heartbeat_max_latency_ms=6_000,
        )
        load = self._load(load=8.5, vcpu=8)
        self.assertEqual(
            rc.classify_window(big_drops, self.policy, system_load=load),
            rc.SLIP_NETWORK,
        )
        # Trivial drop ratio: drops are noise from the AIMD probe; CPU is
        # actually the dominant cause. Hold rate.
        tiny_drops = make_measurement(
            set_rate=4_000_000,
            achieved_pps=3_900_000,
            tx_dropped_delta=10,  # 10 of 117M = 8.5e-8 << threshold
            heartbeat_max_latency_ms=6_000,
        )
        self.assertEqual(
            rc.classify_window(tiny_drops, self.policy, system_load=load),
            rc.SLIP_CPU,
        )

    def test_legacy_classify_when_no_system_load_supplied(self) -> None:
        # Drops + heartbeat, no system_load -> legacy SLIP/SLIP_NETWORK.
        # Ensures bundles without the loadavg reader keep their pre-PR
        # behavior. Asserted both via the constant alias and the new name.
        measurement = make_measurement(
            set_rate=2_000_000,
            achieved_pps=1_900_000,
            tx_dropped_delta=42,
        )
        self.assertEqual(rc.classify_window(measurement, self.policy), rc.SLIP)
        self.assertEqual(rc.classify_window(measurement, self.policy), rc.SLIP_NETWORK)
        self.assertEqual(
            rc.compute_next_rate(self.policy, 2_000_000, measurement),
            1_000_000,
        )

    def test_low_load_with_heartbeat_jitter_classified_network(self) -> None:
        # Load is low but heartbeat slipped — unusual, but legacy behavior
        # treats this as network-style slip (rate-side response). Preserve.
        measurement = make_measurement(
            set_rate=2_000_000,
            achieved_pps=1_900_000,
            heartbeat_max_latency_ms=6_000,
        )
        load = self._load(load=1.0, vcpu=8)  # 0.125 per vcpu
        self.assertEqual(
            rc.classify_window(measurement, self.policy, system_load=load),
            rc.SLIP_NETWORK,
        )

    def test_drops_with_low_load_classified_network(self) -> None:
        # Drops with no CPU pressure: classic NIC saturation, halve.
        measurement = make_measurement(
            set_rate=2_000_000,
            achieved_pps=1_950_000,
            tx_dropped_delta=500,
            heartbeat_max_latency_ms=200,
        )
        load = self._load(load=1.0, vcpu=8)
        self.assertEqual(
            rc.classify_window(measurement, self.policy, system_load=load),
            rc.SLIP_NETWORK,
        )

    def test_clean_window_with_load_info_still_clean(self) -> None:
        # No drops, no jitter, achieved == set, even on a busy box: clean.
        measurement = make_measurement(set_rate=1_000_000, achieved_pps=970_000)
        load = self._load(load=10.0, vcpu=8)  # busy but not slipping
        self.assertEqual(
            rc.classify_window(measurement, self.policy, system_load=load),
            rc.CLEAN,
        )

    def test_invalid_policy_rejects_bad_thresholds(self) -> None:
        with self.assertRaises(ValueError):
            rc.AimdPolicy(cpu_load_threshold=0.0)
        with self.assertRaises(ValueError):
            rc.AimdPolicy(drop_ratio_threshold=-0.1)
        with self.assertRaises(ValueError):
            rc.AimdPolicy(drop_ratio_threshold=1.5)


class SystemLoadReaderTests(unittest.TestCase):
    def test_reads_one_minute_load_from_synthetic_proc(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "loadavg"
            path.write_text("8.50 6.10 5.00 1/3000 1234\n")
            reader = rc.SystemLoadReader(loadavg_path=path, vcpu_count=8)
            load = reader.read()
            self.assertAlmostEqual(load.load_average_1min, 8.5)
            self.assertEqual(load.vcpu_count, 8)
            self.assertAlmostEqual(load.load_average_per_vcpu, 8.5 / 8)

    def test_missing_proc_returns_zero_load(self) -> None:
        reader = rc.SystemLoadReader(loadavg_path="/nonexistent/loadavg", vcpu_count=4)
        load = reader.read()
        self.assertEqual(load.load_average_1min, 0.0)
        self.assertEqual(load.vcpu_count, 4)

    def test_corrupt_proc_returns_zero_load(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            path = Path(tmpdir) / "loadavg"
            path.write_text("not-a-number\n")
            reader = rc.SystemLoadReader(loadavg_path=path, vcpu_count=2)
            self.assertEqual(reader.read().load_average_1min, 0.0)

    def test_per_vcpu_handles_zero_vcpu_gracefully(self) -> None:
        # Defense-in-depth: vcpu_count=0 must not divide-by-zero.
        load = rc.SystemLoad(load_average_1min=4.0, vcpu_count=0)
        self.assertEqual(load.load_average_per_vcpu, 4.0)


class InstanceDefaultsTests(unittest.TestCase):
    def test_lookup_known_classes(self) -> None:
        metal = rc.lookup_instance_defaults("c6in.metal")
        self.assertEqual(metal.starting_rate, 4_000_000)
        self.assertEqual(metal.ceiling, 12_000_000)
        small = rc.lookup_instance_defaults("c6in.xlarge")
        self.assertEqual(small.starting_rate, 500_000)
        self.assertEqual(small.ceiling, 2_000_000)
        m5 = rc.lookup_instance_defaults("m5.xlarge")
        self.assertEqual(m5.starting_rate, 200_000)

    def test_lookup_unknown_returns_none(self) -> None:
        self.assertIsNone(rc.lookup_instance_defaults("alien.42xlarge"))
        self.assertIsNone(rc.lookup_instance_defaults(""))
        self.assertIsNone(rc.lookup_instance_defaults(None))

    def test_apply_fills_unset_knobs_for_metal(self) -> None:
        policy = rc.AimdPolicy()  # defaults
        new_policy, new_starting = rc.apply_instance_defaults(
            policy=policy,
            fallback_rate=500_000,
            instance_type="c6in.metal",
            env={},
        )
        self.assertEqual(new_policy.floor, 1_000_000)
        self.assertEqual(new_policy.ceiling, 12_000_000)
        self.assertEqual(new_starting, 4_000_000)

    def test_apply_preserves_explicit_env_overrides(self) -> None:
        policy = rc.policy_from_env(
            {
                "ANYSCAN_RATE_FLOOR": "200000",
                "ANYSCAN_RATE_CEILING": "5000000",
            }
        )
        new_policy, new_starting = rc.apply_instance_defaults(
            policy=policy,
            fallback_rate=750_000,
            instance_type="c6in.metal",
            env={
                "ANYSCAN_RATE_FLOOR": "200000",
                "ANYSCAN_RATE_CEILING": "5000000",
                "SCANNER_DEFAULT_RATE": "750000",
            },
        )
        self.assertEqual(new_policy.floor, 200_000)
        self.assertEqual(new_policy.ceiling, 5_000_000)
        self.assertEqual(new_starting, 750_000)

    def test_apply_partial_overrides_only_fills_missing(self) -> None:
        # Operator pinned floor but not ceiling: floor stays at env value,
        # ceiling gets the metal default.
        policy = rc.policy_from_env({"ANYSCAN_RATE_FLOOR": "300000"})
        new_policy, new_starting = rc.apply_instance_defaults(
            policy=policy,
            fallback_rate=500_000,
            instance_type="c6in.metal",
            env={"ANYSCAN_RATE_FLOOR": "300000"},
        )
        self.assertEqual(new_policy.floor, 300_000)
        self.assertEqual(new_policy.ceiling, 12_000_000)
        self.assertEqual(new_starting, 4_000_000)

    def test_apply_unknown_instance_type_is_no_op(self) -> None:
        policy = rc.AimdPolicy()
        new_policy, new_starting = rc.apply_instance_defaults(
            policy=policy,
            fallback_rate=500_000,
            instance_type=None,
            env={},
        )
        self.assertEqual(new_policy, policy)
        self.assertEqual(new_starting, 500_000)

    def test_detect_via_dmi_when_product_name_is_real(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            dmi = Path(tmpdir) / "product_name"
            dmi.write_text("c6in.metal\n")
            detected = rc.detect_instance_type(
                env={},
                dmi_path=dmi,
                imds_fetcher=lambda: None,
            )
            self.assertEqual(detected, "c6in.metal")

    def test_detect_skips_dmi_amazon_ec2_falls_back_to_imds(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            dmi = Path(tmpdir) / "product_name"
            dmi.write_text("Amazon EC2\n")
            detected = rc.detect_instance_type(
                env={},
                dmi_path=dmi,
                imds_fetcher=lambda: "c6in.xlarge",
            )
            self.assertEqual(detected, "c6in.xlarge")

    def test_detect_env_override_short_circuits(self) -> None:
        # The IMDS fetcher MUST NOT run when env supplies the type.
        called: list[bool] = []

        def fetcher() -> str:
            called.append(True)
            return "c6in.xlarge"

        detected = rc.detect_instance_type(
            env={"ANYSCAN_INSTANCE_TYPE": "c6in.metal"},
            dmi_path="/nonexistent",
            imds_fetcher=fetcher,
        )
        self.assertEqual(detected, "c6in.metal")
        self.assertEqual(called, [])

    def test_detect_returns_none_when_all_sources_fail(self) -> None:
        detected = rc.detect_instance_type(
            env={},
            dmi_path="/nonexistent/path",
            imds_fetcher=lambda: None,
        )
        self.assertIsNone(detected)

    def test_detect_rejects_garbage_strings(self) -> None:
        # Random product strings (e.g. on-prem hardware) must not match.
        with tempfile.TemporaryDirectory() as tmpdir:
            dmi = Path(tmpdir) / "product_name"
            dmi.write_text("PowerEdge R740xd\n")
            detected = rc.detect_instance_type(
                env={},
                dmi_path=dmi,
                imds_fetcher=lambda: None,
            )
            self.assertIsNone(detected)


class PartialWindowCalibrationTests(unittest.TestCase):
    """Verify calibration is persisted on every clean window AND on terminal."""

    def _stub_runner(
        self,
        rates_then_crash: list[tuple[int, str]],
    ) -> tuple[rc.WindowRunner, list[int]]:
        """Build a runner that emits exactly the given (rate, classification)
        pairs and then raises a mid-window crash on the next call."""

        seen_rates: list[int] = []
        cursor = {"i": 0}

        class ScriptedRunner(rc.WindowRunner):
            def run(self, *, rate, window_seconds, is_first_window):
                seen_rates.append(rate)
                idx = cursor["i"]
                cursor["i"] += 1
                if idx >= len(rates_then_crash):
                    # Simulate scanner crash mid-window.
                    return rc.WindowMeasurement(
                        set_rate=rate,
                        elapsed_seconds=1.0,
                        tx_packets_delta=0,
                        tx_dropped_delta=0,
                        heartbeat_max_latency_ms=0,
                        scanner_finished_naturally=False,
                        scanner_exit_code=139,
                    )
                _expected_rate, kind = rates_then_crash[idx]
                if kind == "clean":
                    return make_measurement(
                        set_rate=rate, achieved_pps=rate * 0.97
                    )
                if kind == "slip":
                    return make_measurement(
                        set_rate=rate,
                        achieved_pps=rate * 0.95,
                        tx_dropped_delta=42,
                    )
                raise AssertionError(f"unknown kind: {kind}")

        return ScriptedRunner(), seen_rates

    def test_persists_on_each_clean_window(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            calib = rc.RateCalibrationStore(Path(tmpdir) / "rate-calibration.json")
            policy = rc.AimdPolicy(window_seconds=30)
            # Clean window 1 at 500k -> bumps to 700k. Clean window 2 at 700k.
            # Then a mid-window crash on window 3.
            runner, seen = self._stub_runner(
                [
                    (500_000, "clean"),
                    (700_000, "clean"),
                ]
            )
            controller = rc.RateController(
                options=rc.ControllerOptions(
                    policy=policy,
                    window_seconds=float(policy.window_seconds),
                    interface="eth0",
                    starting_rate=500_000,
                    calibration=calib,
                ),
                runner=runner,
                log_sink=io.StringIO(),
            )
            with self.assertRaises(rc.ScannerWindowError):
                controller.run()
            self.assertEqual(seen, [500_000, 700_000, 900_000])
            # Calibration must reflect the highest CLEAN rate seen (700_000),
            # NOT 500_000 (older snapshot would persist that) and NOT
            # 900_000 (the crashed rate).
            entry = calib.lookup("eth0")
            self.assertIsNotNone(entry)
            self.assertEqual(entry.learned_rate, 700_000)

    def test_persists_when_crash_strikes_before_any_clean_window(self) -> None:
        # If the scanner crashes on window 1, max_clean_rate is 0 and we
        # must NOT persist anything (regression guard for not stomping
        # the prior calibration with 0).
        with tempfile.TemporaryDirectory() as tmpdir:
            calib = rc.RateCalibrationStore(Path(tmpdir) / "rate-calibration.json")
            calib.store("eth0", 1_500_000, now_iso="2026-04-26T00:00:00Z")
            runner, _ = self._stub_runner([])
            controller = rc.RateController(
                options=rc.ControllerOptions(
                    policy=rc.AimdPolicy(window_seconds=30),
                    window_seconds=30.0,
                    interface="eth0",
                    starting_rate=500_000,
                    calibration=calib,
                ),
                runner=runner,
                log_sink=io.StringIO(),
            )
            with self.assertRaises(rc.ScannerWindowError):
                controller.run()
            entry = calib.lookup("eth0")
            self.assertIsNotNone(entry)
            self.assertEqual(entry.learned_rate, 1_500_000)  # untouched

    def test_persists_when_interrupted_by_systemexit_mid_loop(self) -> None:
        """Signal-handler exit path: SIGTERM/SIGINT in the adapter raises
        SystemExit (handle_termination in vulnscanner-zmap-adapter.py).
        That exception unwinds through SubprocessWindowRunner.run() and
        out of RateController.run()'s while loop. The controller's
        try/finally must still persist max_clean_rate so the calibration
        learned in the windows that DID complete is not lost on a SIGTERM
        from the multi-NIC parent or agentd.
        """

        class InterruptingRunner(rc.WindowRunner):
            def __init__(self) -> None:
                self._idx = 0

            def run(self, *, rate, window_seconds, is_first_window):
                self._idx += 1
                if self._idx == 1:
                    return make_measurement(set_rate=rate, achieved_pps=rate * 0.97)
                if self._idx == 2:
                    return make_measurement(
                        set_rate=rate, achieved_pps=rate * 0.97
                    )
                # Window 3: simulate the adapter's SIGTERM handler which
                # raises SystemExit(128 + signum). Real signal handlers
                # cannot raise during a child.wait() syscall reliably,
                # but the exception path the controller has to survive
                # is identical to the one a Python-level handler would
                # produce.
                raise SystemExit(143)

        with tempfile.TemporaryDirectory() as tmpdir:
            calib = rc.RateCalibrationStore(Path(tmpdir) / "rate-calibration.json")
            policy = rc.AimdPolicy(window_seconds=30)
            controller = rc.RateController(
                options=rc.ControllerOptions(
                    policy=policy,
                    window_seconds=float(policy.window_seconds),
                    interface="eth2",
                    starting_rate=500_000,
                    calibration=calib,
                ),
                runner=InterruptingRunner(),
                log_sink=io.StringIO(),
            )
            with self.assertRaises(SystemExit):
                controller.run()
            entry = calib.lookup("eth2")
            self.assertIsNotNone(entry, "SystemExit must not bypass terminal persist")
            # Highest CLEAN rate observed before the interrupt was 700k
            # (window 2 ran at 700k after window 1's clean@500k bumped it).
            self.assertEqual(entry.learned_rate, 700_000)

    def test_persists_after_natural_finish(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            calib = rc.RateCalibrationStore(Path(tmpdir) / "rate-calibration.json")
            policy = rc.AimdPolicy(window_seconds=30)
            scenario = StubWindow(
                achieved_pps_for_set={500_000: 480_000, 700_000: 680_000},
                drops_for_set={},
                heartbeat_for_set={},
                natural_finish_after_window=2,
            )
            runner = StubRunner(scenario)
            controller = rc.RateController(
                options=rc.ControllerOptions(
                    policy=policy,
                    window_seconds=float(policy.window_seconds),
                    interface="eth1",
                    starting_rate=500_000,
                    calibration=calib,
                ),
                runner=runner,
                log_sink=io.StringIO(),
                max_windows=5,
            )
            controller.run()
            entry = calib.lookup("eth1")
            self.assertIsNotNone(entry)
            # Window 1: clean at 500k. Window 2: clean at 700k AND scanner
            # finished naturally. Highest clean rate = 700k.
            self.assertEqual(entry.learned_rate, 700_000)


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