#!/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: current = get_json(f"{CONTROL}/jobs/{args.job_id}") if current["state"] == "JOB_STATE_STALLED": args.state = "JOB_STATE_STALLED" args.minimum_sequence = 1 command_wait_state(args) return 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())