214 lines
7.4 KiB
Python
Executable File
214 lines
7.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Assertions used by the Phase 8 v2, stall, and capacity scenarios."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
|
|
from complex_jobs import (
|
|
CONTROL,
|
|
get_json,
|
|
post_json,
|
|
qb_json,
|
|
selection_indices,
|
|
)
|
|
|
|
|
|
def command_qb_hash(args: argparse.Namespace) -> None:
|
|
deadline = time.monotonic() + args.timeout
|
|
last = None
|
|
while time.monotonic() < deadline:
|
|
records = qb_json(args.endpoint, "/torrents/info")
|
|
matches = [
|
|
item for item in records
|
|
if item.get("name") == args.name
|
|
]
|
|
last = matches
|
|
if len(matches) == 1:
|
|
value = str(matches[0].get("hash", ""))
|
|
if len(value) not in {40, 64}:
|
|
raise RuntimeError(f"qBittorrent exposed an invalid hash: {value}")
|
|
print(value)
|
|
return
|
|
if len(matches) > 1:
|
|
raise RuntimeError(f"qBittorrent name lookup is ambiguous: {matches}")
|
|
time.sleep(0.25)
|
|
raise RuntimeError(f"qBittorrent torrent did not appear: {last}")
|
|
|
|
|
|
def command_wait_client(args: argparse.Namespace) -> None:
|
|
deadline = time.monotonic() + args.timeout
|
|
last = None
|
|
while time.monotonic() < deadline:
|
|
clients = get_json(f"{CONTROL}/clients")
|
|
matches = [
|
|
item for item in clients
|
|
if item.get("client_id") == args.client_id
|
|
]
|
|
last = matches
|
|
if (
|
|
len(matches) == 1
|
|
and bool(matches[0].get("connected")) == args.connected
|
|
):
|
|
print(json.dumps({
|
|
"client_id": args.client_id,
|
|
"connected": args.connected,
|
|
"generation": matches[0].get("generation"),
|
|
}, sort_keys=True))
|
|
return
|
|
time.sleep(0.25)
|
|
raise RuntimeError(f"client connection state did not converge: {last}")
|
|
|
|
|
|
def command_wait_state(args: argparse.Namespace) -> None:
|
|
deadline = time.monotonic() + args.timeout
|
|
last = None
|
|
while time.monotonic() < deadline:
|
|
last = get_json(f"{CONTROL}/jobs/{args.job_id}")
|
|
if last["state"] == args.state:
|
|
sequence = int(last.get("last_event_sequence", 0))
|
|
if sequence < args.minimum_sequence:
|
|
time.sleep(0.1)
|
|
continue
|
|
progress = (last.get("latest_event") or {}).get("progress") or {}
|
|
overall = float(progress.get("overallFractionComplete", 0))
|
|
if overall < args.minimum_overall:
|
|
raise RuntimeError(
|
|
f"job lost overall progress while {args.state}: {last}"
|
|
)
|
|
print(json.dumps({
|
|
"job_id": args.job_id,
|
|
"last_event_sequence": sequence,
|
|
"overall_fraction_complete": overall,
|
|
"state": last["state"],
|
|
}, sort_keys=True))
|
|
return
|
|
if last["state"] in {
|
|
"JOB_STATE_CANCELLED",
|
|
"JOB_STATE_FAILED",
|
|
"JOB_STATE_SUCCEEDED",
|
|
}:
|
|
raise RuntimeError(
|
|
f"job reached {last['state']} while waiting for {args.state}"
|
|
)
|
|
time.sleep(0.1)
|
|
raise RuntimeError(f"job state did not converge: {last}")
|
|
|
|
|
|
def command_mark_stalled(args: argparse.Namespace) -> None:
|
|
deadline = time.monotonic() + args.timeout
|
|
last = None
|
|
while time.monotonic() < deadline:
|
|
outcomes = post_json(f"{CONTROL}/scheduler/advance", {})
|
|
last = outcomes
|
|
if any(
|
|
item.get("job_id") == args.job_id
|
|
and item.get("action") == "marked_stalled"
|
|
for item in outcomes
|
|
):
|
|
args.state = "JOB_STATE_STALLED"
|
|
args.minimum_sequence = 1
|
|
command_wait_state(args)
|
|
return
|
|
time.sleep(0.25)
|
|
raise RuntimeError(f"scheduler did not mark the job stalled: {last}")
|
|
|
|
|
|
def command_assert_placement(args: argparse.Namespace) -> None:
|
|
expected_files = [
|
|
int(item) for item in args.files.split(",") if item
|
|
]
|
|
placements = get_json(f"{CONTROL}/placements")
|
|
matches = [
|
|
item for item in placements
|
|
if item.get("client_id") == args.client_id
|
|
and item.get("state") == "PLACEMENT_STATE_PRESENT"
|
|
and (
|
|
(args.v1 and item.get("info_hash_v1_hex") == args.v1)
|
|
or (args.v2 and item.get("info_hash_v2_hex") == args.v2)
|
|
)
|
|
]
|
|
if len(matches) != 1:
|
|
raise RuntimeError(f"placement identity is missing or ambiguous: {matches}")
|
|
placement = matches[0]
|
|
actual_files = selection_indices(placement["verified_files"])
|
|
if (
|
|
actual_files != expected_files
|
|
or (args.v1 and placement.get("info_hash_v1_hex") != args.v1)
|
|
or (args.v2 and placement.get("info_hash_v2_hex") != args.v2)
|
|
):
|
|
raise RuntimeError(f"placement identity/selection mismatch: {placement}")
|
|
print(json.dumps({
|
|
"client_id": args.client_id,
|
|
"files": actual_files,
|
|
"info_hash_v1_hex": placement.get("info_hash_v1_hex", ""),
|
|
"info_hash_v2_hex": placement.get("info_hash_v2_hex", ""),
|
|
}, sort_keys=True))
|
|
|
|
|
|
def command_scheduler(args: argparse.Namespace) -> None:
|
|
result = post_json(f"{CONTROL}/scheduler/{args.action}", {})
|
|
expected = args.action == "resume"
|
|
if result.get("automatic_scheduler") is not expected:
|
|
raise RuntimeError(f"scheduler gate did not {args.action}: {result}")
|
|
print(json.dumps(result, sort_keys=True))
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser()
|
|
commands = parser.add_subparsers(dest="command", required=True)
|
|
|
|
qb_hash = commands.add_parser("qb-hash")
|
|
qb_hash.add_argument("--endpoint", required=True)
|
|
qb_hash.add_argument("--name", required=True)
|
|
qb_hash.add_argument("--timeout", type=float, default=30)
|
|
qb_hash.set_defaults(run=command_qb_hash)
|
|
|
|
client = commands.add_parser("wait-client")
|
|
client.add_argument("--client-id", required=True)
|
|
client.add_argument(
|
|
"--connected", action=argparse.BooleanOptionalAction, default=True
|
|
)
|
|
client.add_argument("--timeout", type=float, default=30)
|
|
client.set_defaults(run=command_wait_client)
|
|
|
|
state = commands.add_parser("wait-state")
|
|
state.add_argument("--job-id", required=True)
|
|
state.add_argument("--state", required=True)
|
|
state.add_argument("--minimum-sequence", type=int, default=0)
|
|
state.add_argument("--minimum-overall", type=float, default=0)
|
|
state.add_argument("--timeout", type=float, default=30)
|
|
state.set_defaults(run=command_wait_state)
|
|
|
|
stalled = commands.add_parser("mark-stalled")
|
|
stalled.add_argument("--job-id", required=True)
|
|
stalled.add_argument("--minimum-overall", type=float, default=0.01)
|
|
stalled.add_argument("--timeout", type=float, default=15)
|
|
stalled.set_defaults(run=command_mark_stalled)
|
|
|
|
placement = commands.add_parser("assert-placement")
|
|
placement.add_argument("--client-id", required=True)
|
|
placement.add_argument("--v1", default="")
|
|
placement.add_argument("--v2", default="")
|
|
placement.add_argument("--files", required=True)
|
|
placement.set_defaults(run=command_assert_placement)
|
|
|
|
scheduler = commands.add_parser("scheduler")
|
|
scheduler.add_argument(
|
|
"--action", choices=("pause", "resume"), required=True
|
|
)
|
|
scheduler.set_defaults(run=command_scheduler)
|
|
return parser
|
|
|
|
|
|
def main() -> int:
|
|
args = build_parser().parse_args()
|
|
args.run(args)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|