FabledCurator/tests/test_worker_control.py

"""Changing a lane's slots on a running system (milestone 422 step 2).

The celery control calls are stubbed: what is being tested is the DELTA
ARITHMETIC and the persist/push split, not that celery can resize its own
pool. A test that asserted celery's behaviour would be testing celery.
"""

from __future__ import annotations

import pytest

from backend.app.services import worker_control as wc
from backend.app.services.worker_lanes import LANES_BY_NAME

# --- what inspect reports ----------------------------------------------------


class _Control:
    """Records the control messages that were sent."""

    def __init__(self):
        self.grew: list[tuple[int, list[str]]] = []
        self.shrank: list[tuple[int, list[str]]] = []
        self.added: list[tuple[str, list[str]]] = []
        self.cancelled: list[tuple[str, list[str]]] = []

    def pool_grow(self, n, destination=None):
        self.grew.append((n, destination))

    def pool_shrink(self, n, destination=None):
        self.shrank.append((n, destination))

    def add_consumer(self, queue, destination=None):
        self.added.append((queue, destination))

    def cancel_consumer(self, queue, destination=None):
        self.cancelled.append((queue, destination))


def _stub_live(
    monkeypatch, lane_name, *, pools, present=True, reserved=0, consuming=None,
):
    # `consuming` defaults to the lane's full queue set — i.e. an ENABLED
    # lane. Tests for the disabled case pass an empty set explicitly.
    if consuming is None:
        consuming = set(LANES_BY_NAME[lane_name].queues)
    state = wc.LaneLiveState(
        present=present,
        replicas=len(pools),
        hostnames=sorted(pools),
        pools=dict(pools),
        reserved=reserved,
        consuming=set(consuming),
    )
    monkeypatch.setattr(
        wc, "inspect_lanes_sync",
        lambda: {name: (state if name == lane_name else wc.LaneLiveState())
                 for name in LANES_BY_NAME},
    )
    return state


def _stub_control(monkeypatch):
    control = _Control()

    class _Celery:
        pass

    celery = _Celery()
    celery.control = control
    import sys
    import types
    mod = types.ModuleType("backend.app.celery_app")
    mod.celery = celery
    monkeypatch.setitem(sys.modules, "backend.app.celery_app", mod)
    return control


def test_pool_property_is_max_not_sum():
    """`slots` means the pool size of ONE process, so the aggregate shown to
    the operator is the largest replica — not the total. A sum would report 8
    for two replicas of 4 and invite them to 'reduce it to 4', which would
    halve the lane."""
    state = wc.LaneLiveState(pools={"a": 4, "b": 4})
    assert state.pool == 4


def test_pool_is_none_when_nothing_reported():
    """Unknown, never zero — the distinction step 3's reconcile depends on."""
    assert wc.LaneLiveState(present=True).pool is None


# --- the delta arithmetic ----------------------------------------------------


def test_each_replica_gets_its_own_delta(monkeypatch):
    """The bug this exists to prevent: one delta computed from an aggregate
    and applied to every replica. With replicas at 2 and 6 and a target of 4,
    a shared delta moves both the same way and leaves them at 4 and 8 — or 0
    and 4 — depending on which aggregate was used. Per-replica deltas
    converge both on 4.
    """
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 2, "host-b": 6})

    applied, err = wc.set_lane_slots_sync(LANES_BY_NAME["worker"], 4)

    assert (applied, err) == (True, None)
    assert control.grew == [(2, ["host-a"])]
    assert control.shrank == [(2, ["host-b"])]


def test_a_replica_already_at_the_target_is_sent_nothing(monkeypatch):
    """The fixed point step 3's reconcile needs. An enforcer that re-issues a
    grow of zero every tick never converges and re-does its own work forever
    (lesson #4183)."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 4})

    applied, err = wc.set_lane_slots_sync(LANES_BY_NAME["worker"], 4)

    assert (applied, err) == (True, None)
    assert control.grew == []
    assert control.shrank == []


def test_resizing_an_absent_lane_reports_rather_than_raises(monkeypatch):
    _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={}, present=False)

    applied, err = wc.set_lane_slots_sync(LANES_BY_NAME["worker"], 4)

    assert applied is False
    assert "not running" in err


def test_a_replica_with_no_reported_pool_is_named_not_skipped_silently(
    monkeypatch,
):
    """Resize what can be resized, then say which could not. Silence would
    leave a replica running at a size the UI claims it is not."""
    control = _stub_control(monkeypatch)
    state = _stub_live(monkeypatch, "worker", pools={"host-a": 2})
    state.hostnames = ["host-a", "host-b"]  # b answered inspect, not stats
    state.replicas = 2

    applied, err = wc.set_lane_slots_sync(LANES_BY_NAME["worker"], 4)

    assert control.grew == [(2, ["host-a"])]
    assert applied is False
    assert "host-b" in err


# --- enabling and disabling --------------------------------------------------


def test_disabling_cancels_consumers_rather_than_killing_the_worker(monkeypatch):
    """A cancelled consumer keeps the process alive and answering inspect, so
    a disabled lane stays visible. A killed worker reads as ABSENT, which is
    the same signal as a crash — and milestone 365 exists precisely so those
    two do not look alike."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "ml", pools={"host-a": 1})

    applied, err = wc.set_lane_enabled_sync(LANES_BY_NAME["ml"], False)

    assert (applied, err) == (True, None)
    assert control.cancelled == [("ml", ["host-a"])]
    assert control.added == []


def test_enabling_adds_a_consumer_for_every_queue_in_the_lane(monkeypatch):
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 1})

    wc.set_lane_enabled_sync(LANES_BY_NAME["worker"], True)

    assert [q for q, _ in control.added] == list(LANES_BY_NAME["worker"].queues)


# --- lane identity -----------------------------------------------------------


def test_queue_sets_map_back_to_their_lane_in_any_order():
    """celery does not guarantee the order it lists a worker's queues in, so
    the lookup sorts. Unsorted, a lane would intermittently fail to match and
    read as absent."""
    lane = LANES_BY_NAME["worker"]
    assert wc._lane_for_queues(tuple(reversed(lane.queues))) is lane


def test_an_unknown_queue_set_maps_to_no_lane():
    """A deployment slicing CELERY_QUEUES differently has no lane row to
    control. Honest rather than an error — the roster still reports it."""
    assert wc._lane_for_queues(("something", "else")) is None


# --- the reconcile (step 3) --------------------------------------------------


def test_a_settled_system_sends_no_control_messages(monkeypatch):
    """THE property this sweep lives or dies on. It runs every 5 minutes
    forever, so a converged tick must be silent — one broker round trip and
    nothing else. An enforcer that re-issues a grow of zero churns forever and
    buries a real correction in its own noise (lesson #4183).
    """
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 4})

    result = wc.reconcile_lanes_sync({"worker": (4, True)})

    # EVERY control family, not just the pool ones. The first version of this
    # test asserted only grew/shrank and would have passed while the reconcile
    # re-sent add_consumer for all four queues on every tick — harmless per
    # call, unbounded churn in aggregate, and invisible.
    assert control.grew == []
    assert control.shrank == []
    assert control.added == []
    assert control.cancelled == []
    assert result["changed"] == []


def test_a_worker_back_at_its_env_concurrency_is_corrected(monkeypatch):
    """The failure this exists for: a restart drops the pool to CELERY_
    CONCURRENCY, silently below what the operator set."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 2})

    result = wc.reconcile_lanes_sync({"worker": (6, True)})

    assert control.grew == [(4, ["host-a"])]
    assert result["changed"] == ["worker"]


def test_an_absent_lane_is_skipped_not_corrected(monkeypatch):
    """`present=False` is 'nothing answered', not 'zero slots'. Correcting it
    would be a conclusion drawn from an unswept read (snippet #3969) — and
    there is nothing to send the message to anyway."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={}, present=False)

    result = wc.reconcile_lanes_sync({"worker": (6, True)})

    assert result["skipped"] == ["worker"]
    assert result["changed"] == []
    assert control.grew == []
    assert control.shrank == []


def test_one_lane_failing_does_not_stop_the_others(monkeypatch):
    """A broker blip on one lane must not leave the rest un-reconciled for
    another five minutes."""
    control = _stub_control(monkeypatch)
    present = wc.LaneLiveState(
        present=True, replicas=1, hostnames=["host-a"], pools={"host-a": 1},
    )
    absent = wc.LaneLiveState()
    monkeypatch.setattr(
        wc, "inspect_lanes_sync",
        lambda: {"worker": absent, "scheduler": present,
                 "maintenance_long": absent, "ml": absent},
    )

    result = wc.reconcile_lanes_sync({
        "worker": (4, True), "scheduler": (3, True),
    })

    assert result["skipped"] == ["worker"]
    assert result["changed"] == ["scheduler"]
    assert control.grew == [(2, ["host-a"])]


def test_a_lane_disabled_in_settings_but_still_consuming_is_stopped(monkeypatch):
    """The case that made `consuming` necessary: a lane turned off while its
    worker was down comes back consuming, and must be stopped when it
    returns."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "ml", pools={"host-a": 0}, consuming={"ml"})

    result = wc.reconcile_lanes_sync({"ml": (0, False)})

    assert control.cancelled == [("ml", ["host-a"])]
    assert result["changed"] == ["ml"]


def test_an_already_disabled_lane_is_not_cancelled_again(monkeypatch):
    """The other half of the fixed point. `cancel_consumer` on a queue that is
    not being consumed succeeds and does nothing, so an unconditional
    reconcile would churn here forever with no symptom."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "ml", pools={"host-a": 0}, consuming=set())

    result = wc.reconcile_lanes_sync({"ml": (0, False)})

    assert control.cancelled == []
    assert control.added == []
    assert result["changed"] == []


def test_a_lane_with_no_stored_row_is_left_alone(monkeypatch):
    """A lane in LANES but not in `desired` is one whose row has not been
    seeded. Inventing a target here would let the sweep enforce a number that
    disagrees with the migration it is supposed to be upholding."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 2})

    result = wc.reconcile_lanes_sync({})

    assert result["changed"] == []
    assert control.grew == []


def test_the_reconcile_task_is_registered_and_scheduled():
    """A task name only enters `celery.tasks` when its module is imported, and
    a beat entry naming a task that is not registered fails at tick time
    rather than at import — silently, every five minutes."""
    import backend.app.tasks.maintenance  # noqa: F401
    from backend.app.celery_app import celery

    name = "backend.app.tasks.maintenance.reconcile_worker_lanes"
    assert name in celery.tasks
    scheduled = {e["task"] for e in celery.conf.beat_schedule.values()}
    assert name in scheduled


# --- the autoscaler (step 7) --------------------------------------------------
#
# Every case below fixes the LIVE pool and the stored floor to DIFFERENT
# numbers wherever it can. That is deliberate: the first version of this
# function read the stored value as the current one, and with the two equal
# — which is what a single tick of a settled system looks like — every test
# here still passed while the autoscaler could not grow past floor+1 or shrink
# at all. Equal fixtures cannot see that bug.


def _autoscale_live(
    monkeypatch, *, pools, active, reserved, depth, present=True,
):
    state = wc.LaneLiveState(
        present=present,
        replicas=len(pools),
        hostnames=sorted(pools),
        pools=dict(pools),
        active=active,
        reserved=reserved,
        consuming=set(LANES_BY_NAME["worker"].queues),
    )
    monkeypatch.setattr(
        wc, "inspect_lanes_sync",
        lambda: {name: (state if name == "worker" else wc.LaneLiveState())
                 for name in LANES_BY_NAME},
    )
    monkeypatch.setattr(
        wc, "_queue_depths_sync",
        lambda: {q: depth for lane in LANES_BY_NAME.values() for q in lane.queues},
    )
    return state


def _worker(decisions):
    return next(d for d in decisions if d.lane == "worker")


def test_a_prefetched_backlog_still_triggers_growth(monkeypatch):
    """THE case an LLEN-only implementation misses, and the reason `reserved`
    was plumbed through in step 2. Celery prefetches, so a lane can read queue
    depth 0 while holding thirty tasks in worker memory — an autoscaler
    watching LLEN alone sees an idle system and never grows."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 2}, active=2, reserved=30, depth=0)

    d = _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)}))

    assert d.action == "grew"
    assert d.slots == 3
    assert control.grew == [(1, ["host-a"])]


def test_it_keeps_climbing_past_the_floor_on_later_ticks(monkeypatch):
    """The regression test for reading the stored value as the current one.

    The row still says 2 — the autoscaler never writes it — but the live pool
    is already 5 from earlier ticks. The target must be 6. Computing from the
    stored 2 would propose 3, which resizes nothing (the replica is past it),
    reports success, and pins the lane one slot above its floor forever while
    claiming to grow on every tick."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 5}, active=5, reserved=40, depth=0)

    d = _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)}))

    assert d.slots == 6
    assert control.grew == [(1, ["host-a"])]


def test_backlog_with_free_slots_does_nothing(monkeypatch):
    """Depth alone is not a signal — celery is about to pick those up, and
    growing would add children that idle."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 8}, active=1, reserved=0, depth=50)

    assert _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)})).action == "held"
    assert control.grew == []


def test_saturation_is_measured_against_every_replica(monkeypatch):
    """`active` is summed across replicas and `pool` is one replica's size, so
    comparing them calls two half-busy replicas of 4 saturated at 4 active.
    Capacity is 8 here and 4 tasks are running: half the lane is idle."""
    control = _stub_control(monkeypatch)
    _autoscale_live(
        monkeypatch, pools={"host-a": 4, "host-b": 4},
        active=4, reserved=40, depth=0,
    )

    assert _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)})).action == "held"
    assert control.grew == []


def test_a_saturated_lane_with_no_backlog_does_nothing(monkeypatch):
    """The long-running-task case. One slow task holding every slot with an
    empty queue needs no extra slots — growing does not make it finish sooner,
    which is why the operator's 'runs for x time' idea became a UI warning
    rather than a trigger."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 2}, active=2, reserved=0, depth=0)

    assert _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)})).action == "held"
    assert control.grew == []


def test_growth_stops_at_the_cap_and_says_so(monkeypatch):
    """The autoscaler gets no authority the operator does not already have.
    And it SAYS it is capped — that is the moment they would want to know they
    set one."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 4}, active=4, reserved=99, depth=0)

    d = _worker(wc.autoscale_lanes_sync({"worker": (4, 2, True)}))

    assert d.action == "held"
    assert "cap is 4" in d.reason
    assert control.grew == []


def test_a_cleared_backlog_returns_the_lane_to_the_operator_s_value(monkeypatch):
    """Down toward `configured`, one step at a time."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 5}, active=0, reserved=0, depth=0)

    d = _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)}))

    assert d.action == "shrank"
    assert d.slots == 4
    assert control.shrank == [(1, ["host-a"])]


def test_it_never_shrinks_below_what_the_operator_set(monkeypatch):
    """The stored value is the operator's and the autoscaler must not eat it —
    there would be nothing left to restore to."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 2}, active=0, reserved=0, depth=0)

    assert _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)})).action == "held"
    assert control.shrank == []


def test_hysteresis_keeps_a_middling_backlog_from_flapping(monkeypatch):
    """Between the two thresholds nothing happens. Equal thresholds would grow
    and shrink the lane forever as one task arrives and leaves — lesson
    #4183's churn arriving through a different door."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 5}, active=5, reserved=5, depth=0)

    assert _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)})).action == "held"
    assert control.grew == []
    assert control.shrank == []


def test_a_lane_that_did_not_opt_in_is_never_touched(monkeypatch):
    """Off by default, per lane. This is the only sweep that decides rather
    than obeys, so it acts only where someone said it may — and it does not
    even report on a lane it was not given."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 2}, active=2, reserved=99, depth=0)

    assert wc.autoscale_lanes_sync({"worker": (8, 2, False)}) == []
    assert control.grew == []


def test_an_absent_lane_holds_rather_than_guessing(monkeypatch):
    """`present=False` is 'nothing answered', not 'idle'. Deciding from an
    unswept read is snippet #3969's shape."""
    control = _stub_control(monkeypatch)
    _autoscale_live(
        monkeypatch, pools={}, active=0, reserved=0, depth=0, present=False,
    )

    d = _worker(wc.autoscale_lanes_sync({"worker": (8, 2, True)}))

    assert d.action == "held"
    assert "not answering" in d.reason
    assert control.grew == []


def test_a_settled_lane_sends_no_control_messages(monkeypatch):
    """The fixed point, asserted as a property rather than inferred from the
    reported action: this runs every minute forever, so a settled system has
    to be silent or a real correction drowns in the heartbeat."""
    control = _stub_control(monkeypatch)
    _autoscale_live(monkeypatch, pools={"host-a": 2}, active=1, reserved=0, depth=1)

    for _ in range(5):
        assert _worker(
            wc.autoscale_lanes_sync({"worker": (8, 2, True)})
        ).action == "held"
    assert control.grew == []
    assert control.shrank == []


# --- the two sweeps must not fight -------------------------------------------


def test_the_reconcile_does_not_undo_what_the_autoscaler_added(monkeypatch):
    """Otherwise the two would fight every five minutes: grow, revert, grow,
    revert. For an autoscaling lane the stored slots are a FLOOR."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 6})

    wc.reconcile_lanes_sync({"worker": (2, True)}, frozenset({"worker"}))

    assert control.shrank == []


def test_the_reconcile_still_restores_an_autoscaling_lane_that_fell_below(
    monkeypatch,
):
    """A floor is still a floor. A restart drops the lane to its env value and
    this must bring it back to what the operator set."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 1})

    wc.reconcile_lanes_sync({"worker": (4, True)}, frozenset({"worker"}))

    assert control.grew == [(3, ["host-a"])]


def test_a_non_autoscaling_lane_is_still_driven_down_to_its_value(monkeypatch):
    """The floor applies only to lanes that opted in — otherwise turning
    autoscale off would leave the lane stuck at whatever it had grown to."""
    control = _stub_control(monkeypatch)
    _stub_live(monkeypatch, "worker", pools={"host-a": 6})

    wc.reconcile_lanes_sync({"worker": (2, True)}, frozenset())

    assert control.shrank == [(4, ["host-a"])]


def test_the_autoscale_task_is_registered_and_scheduled():
    import backend.app.tasks.maintenance  # noqa: F401
    from backend.app.celery_app import celery

    name = "backend.app.tasks.maintenance.autoscale_worker_lanes"
    assert name in celery.tasks
    assert name in {e["task"] for e in celery.conf.beat_schedule.values()}