"""Tests on Lab 13's recorded retries run (01_retries.py, openai 3.35.1, October 2026) and the video's rerun of the
rush (200 customers, a test server that lets in 20 calls a second, at most 5 saved up).
Run: pytest test_retry_practice.py"""
from retry_practice import (
    Failure,
    GaveUp,
    TokenBucket,
    backoff,
    calls_to_dead_service,
    chance_through,
    is_retryable,
    retry,
    wilson,
)

# The wave at 15 s in the rush without jitter: when each of 113 calls reached the server, in seconds from the first
# (rounded to 1 ms). The server had been idle for 4 s, so its bucket was full.
WAVE = [
    0.0, 0.0, 0.0, 0.026, 0.185, 0.185, 0.197, 0.197, 0.198, 0.198, 0.198, 0.203, 0.206, 0.207, 0.209, 0.209, 0.211,
    0.211, 0.211, 0.212, 0.212, 0.212, 0.215, 0.215, 0.215, 0.22, 0.22, 0.221, 0.221, 0.222, 0.222, 0.225, 0.226, 0.233,
    0.234, 0.234, 0.234, 0.244, 0.244, 0.246, 0.248, 0.25, 0.25, 0.25, 0.251, 0.251, 0.251, 0.252, 0.256, 0.256, 0.256,
    0.257, 0.257, 0.258, 0.26, 0.261, 0.262, 0.263, 0.263, 0.264, 0.264, 0.265, 0.265, 0.267, 0.267, 0.27, 0.27, 0.271,
    0.272, 0.272, 0.272, 0.274, 0.274, 0.279, 0.28, 0.28, 0.28, 0.283, 0.284, 0.284, 0.284, 0.284, 0.285, 0.285, 0.287,
    0.288, 0.288, 0.288, 0.289, 0.293, 0.293, 0.293, 0.293, 0.294, 0.298, 0.299, 0.303, 0.314, 0.314, 0.315, 0.315,
    0.32, 0.32, 0.32, 0.321, 0.321, 0.321, 0.322, 0.324, 0.325, 0.349, 0.349, 0.371,
]


def test_which_failures_to_retry():
    for status in (429, 500, 502, 503, 504, 408, "timeout", "connection"):
        assert is_retryable(status), status
    for status in (400, 401, 403, 404, 422):
        assert not is_retryable(status), status


def test_backoff_doubles_then_caps_and_jitter_spreads_it():
    assert [backoff(n) for n in range(7)] == [0.5, 1.0, 2.0, 4.0, 8.0, 8.0, 8.0]
    assert backoff(2, base=0.25, cap=4) == 1.0
    assert backoff(3, jitter="full", u=0.0) == 0.0 and backoff(3, jitter="full", u=0.25) == 1.0
    assert backoff(3, jitter="equal", u=0.0) == 2.0 and backoff(3, jitter="equal", u=0.5) == 3.0
    assert backoff(9, jitter="full", u=0.999) < 8.0


def test_coin_flip_math_matches_the_run():
    # 50% of calls failing at random, 100 customers each: 57 / 87 / 96 got through with 1 / 3 / 5 tries.
    for tries, got in ((1, 57), (3, 87), (5, 96)):
        lo, hi = wilson(got, 100)
        assert lo < chance_through(0.5, tries) < hi, (tries, got, lo, hi)
    assert round(chance_through(0.5, 5), 4) == 0.9688  # 5 fails in a row: 1 in 32
    lo, hi = wilson(96, 100)
    assert (round(lo, 2), round(hi, 2)) == (0.9, 0.98)
    assert chance_through(0.1, 1) == 0.9


def test_stacked_retries_multiply():
    assert calls_to_dead_service([5, 3]) == 15  # our 5 tries around the SDK's 3: measured 15 calls
    assert calls_to_dead_service([5, 1]) == 5  # SDK retries off: measured 5
    assert calls_to_dead_service([3, 3, 3]) == 27  # code, gateway and SDK
    assert calls_to_dead_service([1]) == 1


def test_retry_loop():
    # Two 503s, then an answer: two waits, from the backoff.
    script = [Failure(503), Failure(503), "answer"]
    slept = []

    def call():
        r = script.pop(0)
        if isinstance(r, Exception):
            raise r
        return r

    assert retry(call, attempts=5, sleep=slept.append) == "answer"
    assert slept == [0.5, 1.0]

    # A 400 is never retried: raised at once, no waits.
    calls = []

    def bad():
        calls.append(1)
        raise Failure(400)

    slept = []
    try:
        retry(bad, attempts=5, sleep=slept.append)
        raise AssertionError("a 400 should be raised")
    except Failure as e:
        assert e.status == 400
    assert calls == [1] and slept == []

    # Retry-After is a floor: wait at least what the server says.
    script = [Failure(429, retry_after=2.0), Failure(429, retry_after=1.0), "answer"]
    slept = []
    assert retry(call, attempts=5, sleep=slept.append) == "answer"
    assert slept == [2.0, 1.0]

    # A dead service: 5 calls, 4 waits (none after the last), then GaveUp.
    calls, slept = [], []

    def dead():
        calls.append(1)
        raise Failure(503)

    try:
        retry(dead, attempts=5, sleep=slept.append)
        raise AssertionError("should give up")
    except GaveUp as e:
        assert e.attempts == 5
    assert len(calls) == 5 and slept == [0.5, 1.0, 2.0, 4.0]


def test_the_wave_lets_in_twelve():
    bucket = TokenBucket(rate=20, burst=5)
    let_in = [bucket.take(t) for t in WAVE]
    assert sum(let_in) == 12 and len(let_in) == 113
    assert let_in[:5] == [True] * 5  # the 5 saved-up tokens go first
    # Quiet seconds don't pile up: after a long idle gap, still at most 5 at once.
    bucket = TokenBucket(rate=20, burst=5)
    assert sum(bucket.take(100.0) for _ in range(20)) == 5