devplacepy/tests/unit/services/nodes/scheduler.py

# retoor <retoor@molodetz.nl>

import random
from collections import Counter
from datetime import datetime, timezone

import pytest

from devplacepy.services.nodes.circuit import CircuitBreaker
from devplacepy.services.nodes.policy import DestinationPolicy
from devplacepy.services.nodes.registry import NodeCounters, NodeRegistry
from devplacepy.services.nodes.scheduler import MEGABYTE, NodeScheduler, in_active_hours
from devplacepy.services.nodes.session import Inflight, NodeSession
from devplacepy.services.nodes.spec import SPEC
from tests.unit.services.nodes.session import FakeSocket

LIMITS = SPEC.limits
TARGET = DestinationPolicy().check("https://www.reddit.com/r/python.json")


class Clock:
    def __init__(self):
        self.now = 1000.0

    def __call__(self):
        return self.now


def settings(**overrides):
    base = {
        "enabled": True,
        "bandwidth_mb_day": 100,
        "max_concurrent": 2,
        "active_from": "",
        "active_to": "",
        "timezone": "UTC",
        "allowed_platforms": ["reddit"],
        "pause_on_battery": True,
        "public_display": False,
    }
    base.update(overrides)
    return base


def node(uid="d1", user="u1", modes=("fetch",), **overrides):
    session = NodeSession(FakeSocket(), {"uid": uid, "user_uid": user, "kind": "native", "asn_class": "residential"})
    session.hello = {"modes": list(modes)}
    session.settings = settings(**overrides)
    session.roll_day()
    return session


@pytest.fixture
def world(local_db):
    clock = Clock()
    registry = NodeRegistry(NodeCounters())
    breaker = CircuitBreaker(clock=clock)
    wall = {"now": datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)}
    scheduler = NodeScheduler(registry, breaker, clock, lambda: wall["now"], random.Random(1))
    scheduler.wall = wall
    scheduler.clock = clock
    return scheduler


def add(world, **kwargs):
    session = node(**kwargs)
    world.registry.admit(session)
    return session


def test_a_healthy_node_is_picked(world):
    session = add(world)
    assert world.pick(TARGET) is session


def test_no_nodes_means_no_pick(world):
    assert world.pick(TARGET) is None


def test_tunnel_only_jobs_are_not_dispatchable_yet(world):
    add(world)
    assert world.pick(TARGET, "tunnel") is None
    assert world.pick(TARGET, "fetch") is not None
    assert world.pick(TARGET, "any") is not None


def test_a_node_must_speak_the_transport(world):
    add(world, modes=("other",))
    assert world.pick(TARGET) is None


def test_a_tunnel_node_serves_tunnel_and_any_but_not_fetch(world):
    session = add(world, modes=("tunnel",))
    assert world.pick(TARGET, "tunnel") is session
    assert world.pick(TARGET, "any") is session
    assert world.pick(TARGET, "fetch") is None
    assert world.transport_of(session, TARGET, "any") == "tunnel"


def test_any_prefers_nodes_that_can_tunnel(world):
    fetcher = add(world, uid="f1", user="u1", modes=("fetch",))
    tunneler = add(world, uid="t1", user="u2", modes=("fetch", "tunnel"))
    assert world.candidates(TARGET, "any") == [tunneler]
    assert world.transport_of(fetcher, TARGET, "any") == "fetch"


def test_platform_opt_out_is_honoured(world):
    add(world, allowed_platforms=[])
    assert world.pick(TARGET) is None


def test_a_node_never_serves_its_own_owner(world):
    add(world, user="alice")
    assert world.pick(TARGET, caller_owner="alice") is None
    assert world.pick(TARGET, caller_owner="bob") is not None


def test_concurrency_is_capped_by_the_lower_of_server_and_donor(world):
    session = add(world, max_concurrent=1)
    session.inflight["j"] = Inflight("j", "h", "reddit.com", "GET", "u", 0, 0)
    assert world.pick(TARGET) is None
    session.inflight.clear()
    assert session.max_concurrent == 1
    session.settings["max_concurrent"] = 99
    assert session.max_concurrent == LIMITS.max_concurrent_per_node
    session.hello["limits"] = {"max_concurrent": 1}
    assert session.max_concurrent == 1


def test_the_donor_bandwidth_cap_stops_assignment(world):
    session = add(world, bandwidth_mb_day=1)
    session.used_today = MEGABYTE - 1
    assert world.pick(TARGET) is session
    session.used_today = MEGABYTE
    assert world.pick(TARGET) is None


def test_a_new_day_resets_the_used_bytes(world):
    session = add(world, bandwidth_mb_day=1)
    session.used_today = MEGABYTE
    session.used_day = "2000-01-01"
    assert world.pick(TARGET) is session
    assert session.used_today == 0


def test_a_node_waits_between_requests_to_the_same_destination(world):
    session = add(world)
    world.mark_started(session, TARGET)
    world.clock.now += LIMITS.min_seconds_per_node_domain - 0.1
    assert world.eligible(TARGET) == []
    world.clock.now += 0.2
    world.clock.now += 1.0 / TARGET.destination.max_rps
    assert world.eligible(TARGET) == [session]


def test_the_destination_has_a_global_rate_ceiling(world):
    first, second = add(world, uid="a", user="ua"), add(world, uid="b", user="ub")
    world.mark_started(first, TARGET)
    world.clock.now += 1.0 / TARGET.destination.max_rps - 0.01
    assert world.pick(TARGET) is None
    world.clock.now += 0.02
    assert world.pick(TARGET) in (first, second)


def test_a_paused_node_is_skipped_until_the_pause_ends(world):
    session = add(world)
    session.paused_until = world.clock.now + 5
    assert world.pick(TARGET) is None
    world.clock.now += 5
    assert world.pick(TARGET) is session


def test_an_open_node_circuit_removes_only_that_node(world):
    bad, good = add(world, uid="bad", user="u1"), add(world, uid="good", user="u2")
    for _ in range(LIMITS.breaker_failures):
        world.breaker.failure(("node", "bad"))
    assert world.eligible(TARGET) == [good]
    world.clock.now += LIMITS.breaker_open_seconds
    assert set(world.eligible(TARGET)) == {bad, good}


def test_an_open_destination_circuit_stops_everything(world):
    add(world)
    for _ in range(LIMITS.breaker_failures):
        world.breaker.failure(("destination", TARGET.destination.suffix))
    assert world.pick(TARGET) is None


def test_least_loaded_then_least_recently_used_wins(world):
    busy, idle_old, idle_new = add(world, uid="a", user="1"), add(world, uid="b", user="2"), add(world, uid="c", user="3")
    busy.inflight["j"] = Inflight("j", "h", "reddit.com", "GET", "u", 0, 0)
    idle_old.last_assigned = 10
    idle_new.last_assigned = 20
    assert world.pick(TARGET) is idle_old


def test_selection_is_fair_across_equal_nodes(world):
    sessions = [add(world, uid=f"d{i}", user=f"u{i}") for i in range(5)]
    picks = Counter()
    for _ in range(500):
        chosen = world.pick(TARGET)
        picks[chosen.device_uid] += 1
        chosen.last_assigned = world.clock.now
        world.clock.now += 0.001
    assert set(picks) == {item.device_uid for item in sessions}
    assert max(picks.values()) - min(picks.values()) <= 1


@pytest.mark.parametrize(
    "start,end,hour,minute,expected",
    [
        ("", "", 3, 0, True),
        ("09:00", "17:00", 9, 0, True),
        ("09:00", "17:00", 16, 59, True),
        ("09:00", "17:00", 17, 0, False),
        ("09:00", "17:00", 8, 59, False),
        ("22:00", "06:00", 23, 0, True),
        ("22:00", "06:00", 5, 59, True),
        ("22:00", "06:00", 6, 0, False),
        ("22:00", "06:00", 12, 0, False),
        ("09:00", "09:00", 3, 0, True),
        ("bogus", "17:00", 3, 0, True),
    ],
)
def test_active_hours_windows(start, end, hour, minute, expected):
    moment = datetime(2026, 1, 1, hour, minute, tzinfo=timezone.utc)
    assert in_active_hours({"active_from": start, "active_to": end, "timezone": "UTC"}, moment) is expected


def test_active_hours_use_the_donors_timezone():
    moment = datetime(2026, 1, 1, 22, 0, tzinfo=timezone.utc)
    amsterdam = {"active_from": "09:00", "active_to": "17:00", "timezone": "Europe/Amsterdam"}
    assert in_active_hours(amsterdam, moment) is False
    morning = datetime(2026, 1, 1, 8, 30, tzinfo=timezone.utc)
    assert in_active_hours(amsterdam, morning) is True
    broken = {"active_from": "09:00", "active_to": "17:00", "timezone": "Mars/Base"}
    assert in_active_hours(broken, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) is True


def test_a_node_outside_its_hours_is_skipped(world):
    add(world, active_from="22:00", active_to="23:00")
    assert world.pick(TARGET) is None
    world.wall["now"] = datetime(2026, 1, 1, 22, 30, tzinfo=timezone.utc)
    assert world.pick(TARGET) is not None


def test_mark_started_records_both_clocks(world):
    session = add(world)
    world.mark_started(session, TARGET)
    assert session.host_started[TARGET.destination.suffix] == world.clock.now
    assert session.last_assigned == world.clock.now
    assert world._destination_started[TARGET.destination.suffix] == world.clock.now


def test_the_per_destination_maps_stay_bounded_by_the_allowlist(world):
    session = add(world)
    policy = DestinationPolicy()
    for index in range(200):
        world.mark_started(session, policy.check(f"https://sub{index}.reddit.com/"))
    assert len(session.host_started) == 1
    assert len(world._destination_started) == 1


def test_window_checks_are_cached_per_minute_and_still_correct(world):
    from devplacepy.services.nodes.scheduler import _window_open

    _window_open.cache_clear()
    window = {"active_from": "09:00", "active_to": "17:00", "timezone": "UTC"}
    noon = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
    for _ in range(100):
        assert in_active_hours(window, noon) is True
    assert _window_open.cache_info().misses == 1
    night = datetime(2026, 1, 1, 3, 0, tzinfo=timezone.utc)
    assert in_active_hours(window, night) is False


def test_choose_applies_the_owner_exclusion_to_a_shared_candidate_list(world):
    mine, theirs = add(world, uid="a", user="alice"), add(world, uid="b", user="bob")
    found = world.candidates(TARGET)
    assert world.choose(found, "alice") is theirs
    assert world.choose(found, "bob") is mine
    assert world.choose([mine], "alice") is None
    assert world.choose([], "") is None


def test_a_datacenter_donor_is_not_scheduled(world):
    session = add(world)
    session.asn_class = "datacenter"
    assert world.pick(TARGET) is None


def test_an_unknown_class_is_refused_unless_the_site_allows_it(world, local_db):
    from devplacepy.database import set_setting

    session = add(world)
    session.asn_class = "unknown"
    assert world.pick(TARGET) is None
    set_setting("nodes_allow_unknown_class", "1")
    try:
        assert world.pick(TARGET) is session
        session.asn_class = "datacenter"
        assert world.pick(TARGET) is None
    finally:
        set_setting("nodes_allow_unknown_class", "0")


def test_mobile_and_residential_are_eligible(world):
    for index, asn_class in enumerate(("mobile", "residential")):
        session = add(world, uid=f"c{index}", user=f"cu{index}")
        session.asn_class = asn_class
    assert len(world.candidates(TARGET)) == 2


def test_an_admin_override_beats_the_class(world):
    session = add(world)
    session.asn_class = "datacenter"
    session.class_override = "allow"
    assert world.pick(TARGET) is session
    session.class_override = ""
    assert world.pick(TARGET) is None


def test_class_refusals_are_counted(world):
    session = add(world)
    session.asn_class = "datacenter"
    before = world.class_refusals
    world.candidates(TARGET)
    assert world.class_refusals > before