devplacepy/tests/api/nodes/jobs.py

# retoor <retoor@molodetz.nl>

import base64
import json
import time

import pytest
from starlette.websockets import WebSocketDisconnect

from devplacepy.database import get_table, set_setting
from devplacepy.database import nodes as store
from devplacepy.services.manager import service_manager
from devplacepy.services.nodes.broker import JobBroker
from devplacepy.services.nodes.policy import DestinationPolicy
from devplacepy.services.nodes.spec import CLOSE_PROTOCOL, SPEC
from tests.api.nodes.conftest import make_member
from tests.api.nodes.ws import hello, ticket_for

URL = "https://www.reddit.com/r/python.json"
BROKER = JobBroker()


@pytest.fixture
def relaxed(client):
    service = service_manager.get_service("nodes")
    original = service.scheduler.destination_ready
    service.scheduler.destination_ready = lambda target: True
    get_table("node_requests").delete()
    set_setting("nodes_allow_unknown_class", "1")
    yield service
    set_setting("nodes_allow_unknown_class", "0")
    service.scheduler.destination_ready = original
    get_table("node_requests").delete()


def enqueue(caller="caller-1", kind="user", url=URL):
    target = DestinationPolicy().check(url)
    return BROKER.enqueue(kind, caller, target, "any")


def connect(client, member):
    device_uid, _, ticket = ticket_for(client, member)
    return device_uid, client.websocket_connect(f"/nodes/ws?ticket={ticket}")


def next_job(ws, timeout=5.0):
    deadline = time.time() + timeout
    while time.time() < deadline:
        frame = json.loads(ws.receive_text())
        if frame["type"] == "job":
            return frame
        if frame["type"] == "ping":
            ws.send_text(json.dumps({"type": "pong"}))
    raise AssertionError("no job arrived")


def wait_for(uid, status, timeout=5.0):
    deadline = time.time() + timeout
    while time.time() < deadline:
        row = store.get_request(uid)
        if row["status"] == status:
            return row
        time.sleep(0.05)
    raise AssertionError(f"request stayed {store.get_request(uid)['status']}")


def test_a_queued_request_reaches_the_node_and_the_result_comes_back(client, relaxed, member):
    device_uid, connection = connect(client, member)
    with connection as ws:
        ws.send_text(hello(modes=["fetch"]))
        assert json.loads(ws.receive_text())["type"] == "welcome"
        uid = enqueue()
        job = next_job(ws)
        assert job["job_id"] == uid and job["url"] == URL and job["method"] == "GET"
        assert job["timeout"] == SPEC.limits.job_timeout_seconds
        ws.send_text(
            json.dumps(
                {
                    "type": "job_result",
                    "job_id": uid,
                    "status": 200,
                    "headers": {"Content-Type": "application/json"},
                    "body_b64": base64.b64encode(b'{"ok": true}').decode(),
                }
            )
        )
        wait_for(uid, "done")
    result = BROKER.read(uid)
    assert result.http_status == 200 and result.body == b'{"ok": true}'
    assert result.headers == {"content-type": "application/json"}
    today = store.usage_today(device_uid)
    assert today["requests_ok"] == 1 and today["bytes_down"] == 12
    assert store.get_donor(member["uid"])["requests_ok"] == 1


def test_a_node_that_speaks_no_known_transport_gets_nothing(client, relaxed, member):
    _, connection = connect(client, member)
    with connection as ws:
        ws.send_text(hello(modes=["any"]))
        ws.receive_text()
        uid = enqueue()
        time.sleep(0.7)
        assert store.get_request(uid)["status"] == "queued"


def test_a_donor_never_receives_its_own_request(client, relaxed, member):
    _, connection = connect(client, member)
    with connection as ws:
        ws.send_text(hello(modes=["fetch"]))
        ws.receive_text()
        uid = enqueue(caller=member["uid"])
        time.sleep(0.7)
        assert store.get_request(uid)["status"] == "queued"


def test_a_donor_opt_out_of_the_platform_is_respected_live(client, relaxed, member):
    store.save_node_settings(member["uid"], {"allowed_platforms": []})
    _, connection = connect(client, member)
    with connection as ws:
        ws.send_text(hello(modes=["fetch"]))
        assert json.loads(ws.receive_text())["settings"]["allowed_platforms"] == []
        uid = enqueue()
        time.sleep(0.7)
        assert store.get_request(uid)["status"] == "queued"


def test_a_job_error_marks_the_request_failed(client, relaxed, member):
    _, connection = connect(client, member)
    with connection as ws:
        ws.send_text(hello(modes=["fetch"]))
        ws.receive_text()
        uid = enqueue()
        job = next_job(ws)
        ws.send_text(json.dumps({"type": "job_error", "job_id": job["job_id"], "code": "timeout"}))
        row = wait_for(uid, "failed")
    assert row["error"] == "timeout"


def test_disconnecting_with_a_job_in_flight_fails_it(client, relaxed, member):
    _, connection = connect(client, member)
    with connection as ws:
        ws.send_text(hello(modes=["fetch"]))
        ws.receive_text()
        uid = enqueue()
        next_job(ws)
    row = wait_for(uid, "failed")
    assert row["error"] == "node_disconnected"


def test_a_forged_result_is_an_integrity_failure_and_a_flag(client, relaxed, member):
    _, connection = connect(client, member)
    with connection as ws:
        ws.send_text(hello(modes=["fetch"]))
        ws.receive_text()
        uid = enqueue()
        next_job(ws)
        ws.send_text(
            json.dumps({"type": "job_result", "job_id": uid, "status": 200, "body_b64": "***"})
        )
        row = wait_for(uid, "failed")
    assert row["error"] == "bad_base64"
    assert any(flag["kind"] == "integrity" for flag in store.list_flags())


def test_answering_a_job_that_was_never_assigned_closes_the_socket(client, relaxed, member):
    _, connection = connect(client, member)
    with connection as ws:
        ws.send_text(hello(modes=["fetch"]))
        ws.receive_text()
        ws.send_text(
            json.dumps({"type": "job_result", "job_id": "made-up", "status": 200, "body_b64": ""})
        )
        with pytest.raises(WebSocketDisconnect) as info:
            ws.receive_text()
        assert info.value.code == CLOSE_PROTOCOL


def test_two_donors_split_two_requests(client, relaxed):
    first, second = make_member(), make_member()
    _, one = connect(client, first)
    _, two = connect(client, second)
    with one as ws_one, two as ws_two:
        for ws in (ws_one, ws_two):
            ws.send_text(hello(modes=["fetch"]))
            ws.receive_text()
        a = enqueue(caller="x")
        b = enqueue(caller="y", url="https://old.reddit.com/r/rust.json")
        jobs = {next_job(ws_one)["job_id"], next_job(ws_two)["job_id"]}
        assert jobs == {a, b}