"""Run from a saved .py file, not stdin: spawn must import the main module."""
from concurrent.futures import ProcessPoolExecutor
import json
import multiprocessing as mp
import subprocess
import sys
from count_primes import count_chunk


def main():
    expected = {"limit": 10000, "chunks": 4, "prime_count": 1229}
    for extra in (["--serial"], [], ["--workers", "1"]):
        result = subprocess.run(
            [sys.executable, "count_primes.py", *extra],
            capture_output=True, text=True, timeout=30, check=False,
        )
        assert result.returncode == 0 and not result.stderr, result
        assert json.loads(result.stdout) == expected, result
    for extra in (["--workers", "0"], ["--limit", "1"],
                  ["--chunk-size", "0"], ["--chunk-size", "1"]):
        result = subprocess.run(
            [sys.executable, "count_primes.py", *extra],
            capture_output=True, text=True, timeout=30, check=False,
        )
        assert result.returncode == 2 and result.stdout == "", result
        assert "error:" in result.stderr, result
    with ProcessPoolExecutor(max_workers=1, mp_context=mp.get_context("spawn")) as pool:
        future = pool.submit(count_chunk, (8, 2))
        try:
            future.result(timeout=10)
        except ValueError as exc:
            assert str(exc) == "invalid interval", exc
        else:
            raise AssertionError("worker failure was lost")
    print("PASS: serial/process/one-worker agree; bad arguments rejected; worker error received")


if __name__ == "__main__":
    main()