1
0
forked from BRT/arc
arc/tests/mock_server/inject_random.py

68 lines
2.2 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Pick N random scenario JSONs from a pool and inject them into the mock server.
Usage:
./inject_random.py # 5 random from scenarios/pool/
./inject_random.py --count 3
./inject_random.py --dir scenarios/pool --count 5 --host 127.0.0.1 --port 8888
Each injected task is tagged with a sequential `_tag` field (if present) so
you can match the `postTaskResult` responses back to which pool entry ran.
"""
import argparse
import json
import random
import sys
import urllib.request
from pathlib import Path
def main() -> int:
here = Path(__file__).parent
parser = argparse.ArgumentParser()
parser.add_argument("--dir", default=str(here / "scenarios" / "pool"),
help="Directory with scenario JSON files")
parser.add_argument("--count", type=int, default=5,
help="How many scenarios to pick randomly")
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8888)
parser.add_argument("--seed", type=int, default=None,
help="Optional RNG seed for reproducibility")
args = parser.parse_args()
pool_dir = Path(args.dir)
files = sorted(pool_dir.glob("*.json"))
if not files:
print(f"no scenarios in {pool_dir}", file=sys.stderr)
return 1
if args.count > len(files):
print(f"asked for {args.count} but pool has only {len(files)}",
file=sys.stderr)
return 1
if args.seed is not None:
random.seed(args.seed)
picked = random.sample(files, args.count)
tasks = []
for p in picked:
with p.open() as f:
tasks.append(json.load(f))
print(f" picked: {p.name}")
url = f"http://{args.host}:{args.port}/_admin/inject"
req = urllib.request.Request(
url,
data=json.dumps(tasks).encode(),
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req, timeout=5) as resp:
print(f"injected {len(tasks)} tasks -> {resp.status} {resp.read().decode()}")
return 0
if __name__ == "__main__":
sys.exit(main())