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

# retoor <retoor@molodetz.nl>

import json
from datetime import datetime, timezone

import pytest

from devplacepy.database import get_table, set_setting
from devplacepy.database import nodes as store
from devplacepy.services.nodes.broker import BrokerResult, JobBroker
from devplacepy.services.nodes.cache import ResultCache
from devplacepy.services.nodes.fabric import (
    FabricQuotaExceeded,
    FabricRefused,
    FabricUnavailable,
    NodeFabric,
)
from devplacepy.services.nodes.spec import SPEC
from devplacepy.utils import generate_uid
from tests.conftest import run_async

LIMITS = SPEC.limits
URL = "https://www.reddit.com/r/python.json"


class ScriptedBroker(JobBroker):
    def __init__(self, outcomes):
        self.outcomes = list(outcomes)
        self.enqueued = 0

    def enqueue(self, caller_kind, caller_id, target, mode):
        self.enqueued += 1
        return f"uid-{self.enqueued}"

    async def wait(self, uid, timeout):
        outcome = self.outcomes.pop(0)
        if outcome == "ok":
            return BrokerResult(uid, "done", 200, {"content-type": "text/plain"}, b"payload", False, "")
        return BrokerResult(uid, "failed", 0, {}, b"", False, outcome)


def heartbeat(connected):
    table = get_table("service_state")
    record = {
        "name": "nodes",
        "status": "running",
        "heartbeat": datetime.now(timezone.utc).isoformat(),
        "metrics": json.dumps({"stats": [{"label": "connected", "value": connected}]}),
    }
    existing = table.find_one(name="nodes")
    if existing:
        table.update({**record, "id": existing["id"]}, ["id"])
    else:
        table.insert(record)


@pytest.fixture
def online(local_db):
    set_setting("nodes_enabled", "1")
    set_setting("nodes_daily_quota_user", "1000")
    set_setting("nodes_daily_quota_guest", "0")
    heartbeat(3)
    store.cache_clear()
    yield
    store.cache_clear()
    set_setting("nodes_enabled", "0")
    get_table("service_state").delete(name="nodes")


def fetch(fabric, **kwargs):
    kwargs.setdefault("caller_kind", "user")
    kwargs.setdefault("caller_id", generate_uid())
    return run_async(fabric.fetch(URL, **kwargs))


def test_a_node_fetch_returns_the_body_and_caches_it(online):
    broker = ScriptedBroker(["ok"])
    fabric = NodeFabric(broker=broker)
    first = fetch(fabric, caller_id="c")
    assert (first.status, first.body, first.via) == (200, b"payload", "fetch")
    second = fetch(fabric, caller_id="c")
    assert second.via == "cache" and second.body == b"payload" and second.elapsed_ms == 0
    assert broker.enqueued == 1


def test_use_cache_false_always_asks_a_node(online):
    broker = ScriptedBroker(["ok", "ok"])
    fabric = NodeFabric(broker=broker)
    fetch(fabric, caller_id="c", use_cache=False)
    fetch(fabric, caller_id="c", use_cache=False)
    assert broker.enqueued == 2


def test_policy_refusals_surface_with_their_code(online):
    fabric = NodeFabric(broker=ScriptedBroker([]))
    for url, code in (("https://example.com/", "host_not_allowed"), ("https://reddit.com:22/", "port_not_allowed")):
        with pytest.raises(FabricRefused) as info:
            run_async(fabric.fetch(url, caller_kind="user", caller_id="c"))
        assert info.value.code == code


def test_the_kill_switch_and_missing_nodes_are_unavailable(online):
    fabric = NodeFabric(broker=ScriptedBroker([]))
    heartbeat(0)
    with pytest.raises(FabricUnavailable) as info:
        fetch(fabric)
    assert info.value.code == "no_nodes"
    heartbeat(2)
    set_setting("nodes_enabled", "0")
    with pytest.raises(FabricUnavailable) as info:
        fetch(fabric)
    assert info.value.code == "kill_switch"


def test_a_stale_heartbeat_counts_as_no_nodes(online):
    table = get_table("service_state")
    row = table.find_one(name="nodes")
    table.update({"id": row["id"], "heartbeat": "2000-01-01T00:00:00+00:00"}, ["id"])
    assert NodeFabric(broker=ScriptedBroker([])).available().nodes_online == 0


def test_no_service_row_means_no_nodes(online):
    get_table("service_state").delete(name="nodes")
    assert NodeFabric().available().nodes_online == 0


def test_quota_is_enforced_before_any_node_is_touched(online):
    set_setting("nodes_daily_quota_user", "2")
    broker = ScriptedBroker(["ok"] * 5)
    fabric = NodeFabric(broker=broker)
    caller = generate_uid()
    for _ in range(2):
        fetch(fabric, caller_id=caller, use_cache=False)
    with pytest.raises(FabricQuotaExceeded):
        fetch(fabric, caller_id=caller, use_cache=False)
    assert broker.enqueued == 2


def test_guests_have_no_quota_by_default(online):
    broker = ScriptedBroker(["ok"])
    with pytest.raises(FabricQuotaExceeded):
        fetch(NodeFabric(broker=broker), caller_kind="guest", caller_id="g")
    assert broker.enqueued == 0


def test_a_cache_hit_costs_no_quota(online):
    set_setting("nodes_daily_quota_user", "1")
    fabric = NodeFabric(broker=ScriptedBroker(["ok"]))
    caller = generate_uid()
    fetch(fabric, caller_id=caller)
    for _ in range(5):
        assert fetch(fabric, caller_id=caller).via == "cache"


def test_a_retryable_failure_is_retried_once_on_another_node(online):
    broker = ScriptedBroker(["node_disconnected", "ok"])
    result = fetch(NodeFabric(broker=broker))
    assert result.body == b"payload" and broker.enqueued == 2


def test_the_retry_does_not_spend_a_second_quota_unit(online):
    set_setting("nodes_daily_quota_user", "1")
    caller = generate_uid()
    fetch(NodeFabric(broker=ScriptedBroker(["timeout", "ok"])), caller_id=caller)
    assert store.caller_requests_today("user", caller) == 1


def test_retries_stop_after_the_attempt_cap(online):
    broker = ScriptedBroker(["timeout"] * 5)
    with pytest.raises(FabricUnavailable) as info:
        fetch(NodeFabric(broker=broker))
    assert info.value.code == "timeout" and broker.enqueued == 2


def test_a_non_retryable_failure_is_not_retried(online):
    broker = ScriptedBroker(["bad_status", "ok"])
    with pytest.raises(FabricUnavailable):
        fetch(NodeFabric(broker=broker))
    assert broker.enqueued == 1


def test_a_full_queue_is_a_clean_unavailable(online):
    class Full(ScriptedBroker):
        def enqueue(self, *args):
            raise store.QueueFull("full")

    with pytest.raises(FabricUnavailable) as info:
        fetch(NodeFabric(broker=Full([])))
    assert info.value.code == "queue_full"


def test_the_wait_is_capped_at_the_spec_maximum(online):
    seen = []

    class Spy(ScriptedBroker):
        async def wait(self, uid, timeout):
            seen.append(timeout)
            return await super().wait(uid, timeout)

    fetch(NodeFabric(broker=Spy(["ok"])), timeout=10**6)
    assert seen[0] <= LIMITS.max_wait_seconds


def test_the_cache_ttl_follows_the_admin_setting_up_to_the_ceiling(online):
    cache = ResultCache()
    fabric = NodeFabric(broker=ScriptedBroker(["ok"]), cache=cache)
    set_setting("nodes_cache_ttl_seconds", "0")
    fetch(fabric, caller_id="c")
    assert len(cache) == 0
    set_setting("nodes_cache_ttl_seconds", str(10**9))
    fabric._sync_cache_ttl()
    assert cache.ttl_seconds == LIMITS.cache_ttl_seconds
    set_setting("nodes_cache_ttl_seconds", str(LIMITS.cache_ttl_seconds))


def test_can_serve_needs_policy_switch_and_a_node(online):
    fabric = NodeFabric(broker=ScriptedBroker([]))
    assert fabric.can_serve(URL)
    assert not fabric.can_serve("https://example.com/")
    assert not fabric.can_serve(URL, "POST")
    heartbeat(0)
    assert not fabric.can_serve(URL)


def test_the_real_broker_times_out_a_request_nobody_picks_up(online):
    get_table("node_requests").delete()
    fabric = NodeFabric()
    with pytest.raises(FabricUnavailable) as info:
        fetch(fabric, timeout=0.3)
    assert info.value.code in ("caller_timeout", "failed")
    rows = list(get_table("node_requests").find())
    assert rows and all(row["status"] == "failed" for row in rows)
    get_table("node_requests").delete()