captchaforge 0.2.40

//! Unit tests for [`super`] (the bandit solver-selection policy).

use super::*;
use rand::SeedableRng;

fn methods() -> Vec<SolveMethod> {
    vec![
        SolveMethod::BehavioralBypass,
        SolveMethod::VisionLLM,
        SolveMethod::AudioBypass,
        SolveMethod::CrowdSourced,
        SolveMethod::AutoPass,
    ]
}

// ─── Basic semantics ────────────────────────────────────────────

#[test]
fn arm_new_starts_with_uniform_prior() {
    let arm = Arm::new("X".into(), "Y".into());
    assert_eq!(arm.alpha, 1.0);
    assert_eq!(arm.beta, 1.0);
    assert!((arm.posterior_mean() - 0.5).abs() < 0.001);
}

#[test]
fn arm_observe_success_increments_alpha() {
    let mut arm = Arm::new("X".into(), "Y".into());
    arm.observe(true);
    assert_eq!(arm.alpha, 2.0);
    assert_eq!(arm.beta, 1.0);
}

#[test]
fn arm_observe_failure_increments_beta() {
    let mut arm = Arm::new("X".into(), "Y".into());
    arm.observe(false);
    assert_eq!(arm.alpha, 1.0);
    assert_eq!(arm.beta, 2.0);
}

#[test]
fn arm_posterior_mean_tracks_success_rate() {
    let mut arm = Arm::new("X".into(), "Y".into());
    for _ in 0..9 {
        arm.observe(true);
    }
    arm.observe(false);
    // 10 successes (α=10), 2 failures (β=2) → mean = 10/12 = 0.833
    let mean = arm.posterior_mean();
    assert!((mean - 10.0 / 12.0).abs() < 0.001, "got {mean}");
}

// ─── Choose dispatching ─────────────────────────────────────────

#[test]
fn choose_returns_none_on_empty_candidates() {
    let b = Bandit::new();
    let r = b.choose(&CaptchaType::CloudflareTurnstile, &[]);
    assert!(r.is_none());
}

#[test]
fn choose_returns_some_when_candidates_present() {
    let b = Bandit::new();
    let r = b.choose(&CaptchaType::CloudflareTurnstile, &methods());
    assert!(r.is_some());
}

#[test]
fn observe_increments_arm_count() {
    let b = Bandit::new();
    b.observe(
        &CaptchaType::CloudflareTurnstile,
        &SolveMethod::AutoPass,
        true,
    );
    assert_eq!(b.len(), 1);
    b.observe(
        &CaptchaType::CloudflareTurnstile,
        &SolveMethod::VisionLLM,
        false,
    );
    assert_eq!(b.len(), 2);
}

// ─── Persistence ───────────────────────────────────────────────

#[test]
fn save_and_load_round_trips_arms() {
    let b = Bandit::new();
    for _ in 0..5 {
        b.observe(
            &CaptchaType::CloudflareTurnstile,
            &SolveMethod::AutoPass,
            true,
        );
    }
    for _ in 0..2 {
        b.observe(
            &CaptchaType::CloudflareTurnstile,
            &SolveMethod::VisionLLM,
            false,
        );
    }
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("bandit.json");
    b.save_to(&path).unwrap();

    let b2 = Bandit::new();
    let loaded = b2.load_from(&path).unwrap();
    assert_eq!(loaded, 2);
    assert_eq!(b2.len(), 2);
}

// ─── Sampling ──────────────────────────────────────────────────

#[test]
fn sample_beta_in_unit_interval() {
    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
    for _ in 0..1000 {
        let x = sample_beta(2.0, 3.0, &mut rng);
        assert!((0.0..=1.0).contains(&x), "sample {x} out of [0,1]");
    }
}

#[test]
fn sample_beta_mean_matches_alpha_over_alpha_plus_beta() {
    let mut rng = rand::rngs::StdRng::seed_from_u64(1234);
    let (alpha, beta) = (8.0, 2.0);
    let mut sum = 0.0;
    let n = 5000;
    for _ in 0..n {
        sum += sample_beta(alpha, beta, &mut rng);
    }
    let mean = sum / (n as f64);
    let expected = alpha / (alpha + beta);
    assert!(
        (mean - expected).abs() < 0.02,
        "empirical mean {mean} != expected {expected}"
    );
}

// ─── Convergence ───────────────────────────────────────────────

/// 5-arm Bernoulli problem with true means [0.9, 0.6, 0.4, 0.2, 0.1].
/// After 1000 rounds Thompson sampling must concentrate >70% of
/// pulls on the best arm.
#[test]
fn convergence_thompson_concentrates_on_best_arm() {
    let true_p = [0.9, 0.6, 0.4, 0.2, 0.1];
    let methods_local = vec![
        SolveMethod::AutoPass,
        SolveMethod::BehavioralBypass,
        SolveMethod::VisionLLM,
        SolveMethod::AudioBypass,
        SolveMethod::ThirdPartyService,
    ];
    let b = Bandit::new();
    let ct = CaptchaType::CloudflareTurnstile;
    let mut rng = rand::rngs::StdRng::seed_from_u64(0xC0DE);

    let mut counts = [0usize; 5];
    for _ in 0..1000 {
        let chosen = b.choose(&ct, &methods_local).unwrap();
        let idx = methods_local.iter().position(|m| m == &chosen).unwrap();
        counts[idx] += 1;
        let succeeded: bool = rng.gen_bool(true_p[idx]);
        b.observe(&ct, &chosen, succeeded);
    }
    // Arm 0 should dominate.
    let best_pulls = counts[0];
    let total = counts.iter().sum::<usize>();
    let frac = (best_pulls as f64) / (total as f64);
    assert!(
        frac >= 0.55,
        "Thompson didn't concentrate; arm0 got {best_pulls}/{total} = {frac:.3}"
    );
}

/// Total reward over a 1000-round simulation should be within
/// 25% of the oracle (always pulling the best arm).
#[test]
fn convergence_regret_bounded() {
    let true_p = [0.9, 0.5, 0.3, 0.1, 0.05];
    let methods_local = vec![
        SolveMethod::AutoPass,
        SolveMethod::BehavioralBypass,
        SolveMethod::VisionLLM,
        SolveMethod::AudioBypass,
        SolveMethod::ThirdPartyService,
    ];
    let b = Bandit::new();
    let ct = CaptchaType::CloudflareTurnstile;
    let mut rng = rand::rngs::StdRng::seed_from_u64(0xC0FFEE);

    let mut bandit_reward = 0;
    for _ in 0..1000 {
        let chosen = b.choose(&ct, &methods_local).unwrap();
        let idx = methods_local.iter().position(|m| m == &chosen).unwrap();
        let succeeded: bool = rng.gen_bool(true_p[idx]);
        if succeeded {
            bandit_reward += 1;
        }
        b.observe(&ct, &chosen, succeeded);
    }
    let oracle_expected_reward = 1000.0 * true_p[0]; // = 900
    let ratio = (bandit_reward as f64) / oracle_expected_reward;
    assert!(
        ratio >= 0.65,
        "bandit reward {bandit_reward} fell below 65% of oracle {oracle_expected_reward}"
    );
}

/// Scale test: 50k rounds of Bernoulli 5-arm problem. Empirical
/// regret O(√T log T) should give total regret ≤ ~30% of oracle.
#[test]
fn scale_50k_rounds_thompson_competitive() {
    let true_p = [0.85, 0.6, 0.4, 0.2, 0.1];
    let methods_local = vec![
        SolveMethod::AutoPass,
        SolveMethod::BehavioralBypass,
        SolveMethod::VisionLLM,
        SolveMethod::AudioBypass,
        SolveMethod::ThirdPartyService,
    ];
    let b = Bandit::new();
    let ct = CaptchaType::CloudflareTurnstile;
    let mut rng = rand::rngs::StdRng::seed_from_u64(0xDEAD_BEEF);

    let mut bandit_reward = 0;
    for _ in 0..50_000 {
        let chosen = b.choose(&ct, &methods_local).unwrap();
        let idx = methods_local.iter().position(|m| m == &chosen).unwrap();
        let succeeded: bool = rng.gen_bool(true_p[idx]);
        if succeeded {
            bandit_reward += 1;
        }
        b.observe(&ct, &chosen, succeeded);
    }
    let oracle_expected = 50_000.0 * true_p[0];
    let ratio = (bandit_reward as f64) / oracle_expected;
    assert!(
        ratio >= 0.75,
        "50k rounds: bandit reward {} below 75% of oracle {}; ratio={:.3}",
        bandit_reward,
        oracle_expected,
        ratio
    );
}

/// Drift test: best arm changes after 5k rounds. Bandit with
/// decay should re-converge faster than without.
#[test]
fn drift_decay_helps_recovery_after_arm_change() {
    let methods_local = vec![
        SolveMethod::AutoPass,
        SolveMethod::BehavioralBypass,
        SolveMethod::VisionLLM,
    ];
    let mut rng = rand::rngs::StdRng::seed_from_u64(0x1234_5678);

    // Without decay.
    let b_static = Bandit::new();
    let ct = CaptchaType::CloudflareTurnstile;
    let mut switch_recovery_static = 0;
    for round in 0..10_000 {
        let chosen = b_static.choose(&ct, &methods_local).unwrap();
        let idx = methods_local.iter().position(|m| m == &chosen).unwrap();
        // First half: arm 0 wins. Second half: arm 2 wins.
        let p = if round < 5_000 {
            [0.9, 0.4, 0.1][idx]
        } else {
            [0.1, 0.4, 0.9][idx]
        };
        let succeeded: bool = rng.gen_bool(p);
        if round >= 5_000 && chosen == SolveMethod::VisionLLM && succeeded {
            switch_recovery_static += 1;
        }
        b_static.observe(&ct, &chosen, succeeded);
    }

    // With decay.
    let b_decay = Bandit::new().with_decay(0.995);
    let mut switch_recovery_decay = 0;
    for round in 0..10_000 {
        let chosen = b_decay.choose(&ct, &methods_local).unwrap();
        let idx = methods_local.iter().position(|m| m == &chosen).unwrap();
        let p = if round < 5_000 {
            [0.9, 0.4, 0.1][idx]
        } else {
            [0.1, 0.4, 0.9][idx]
        };
        let succeeded: bool = rng.gen_bool(p);
        if round >= 5_000 && chosen == SolveMethod::VisionLLM && succeeded {
            switch_recovery_decay += 1;
        }
        b_decay.observe(&ct, &chosen, succeeded);
    }

    assert!(
        switch_recovery_decay >= switch_recovery_static,
        "decay variant ({switch_recovery_decay}) didn't help vs static ({switch_recovery_static})"
    );
}

// ─── Property tests ────────────────────────────────────────────

proptest::proptest! {
    #![proptest_config(proptest::test_runner::Config {
        cases: 10_000, .. proptest::test_runner::Config::default()
    })]

    #[test]
    fn prop_choose_returns_one_of_candidates(
        n in 1usize..6,
    ) {
        let b = Bandit::new();
        let cand: Vec<SolveMethod> = methods().into_iter().take(n).collect();
        let r = b.choose(&CaptchaType::CloudflareTurnstile, &cand).unwrap();
        assert!(cand.contains(&r));
    }

    #[test]
    fn prop_observe_alpha_beta_strictly_monotonic(
        successes in 0u32..200,
        failures in 0u32..200,
    ) {
        let b = Bandit::new();
        let ct = CaptchaType::CloudflareTurnstile;
        let sm = SolveMethod::AutoPass;
        for _ in 0..successes {
            b.observe(&ct, &sm, true);
        }
        for _ in 0..failures {
            b.observe(&ct, &sm, false);
        }
        let snap = b.snapshot();
        let arm = snap.iter().find(|a| a.solver_method == "AutoPass");
        if successes > 0 || failures > 0 {
            let arm = arm.unwrap();
            assert!(arm.alpha >= 1.0);
            assert!(arm.beta >= 1.0);
            // α grew by exactly `successes`, β by exactly `failures`.
            assert!((arm.alpha - (1.0 + successes as f64)).abs() < 0.001);
            assert!((arm.beta - (1.0 + failures as f64)).abs() < 0.001);
        }
    }

    #[test]
    fn prop_sample_beta_in_unit_interval(
        alpha in 0.1f64..100.0,
        beta in 0.1f64..100.0,
    ) {
        let mut rng = rand::thread_rng();
        for _ in 0..10 {
            let x = sample_beta(alpha, beta, &mut rng);
            assert!((0.0..=1.0).contains(&x));
        }
    }

    #[test]
    fn prop_save_load_idempotent(
        successes in 0u32..50,
        failures in 0u32..50,
    ) {
        let b = Bandit::new();
        let ct = CaptchaType::CloudflareTurnstile;
        let sm = SolveMethod::AutoPass;
        for _ in 0..successes {
            b.observe(&ct, &sm, true);
        }
        for _ in 0..failures {
            b.observe(&ct, &sm, false);
        }
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("bandit.json");
        b.save_to(&path).unwrap();
        let b2 = Bandit::new();
        b2.load_from(&path).unwrap();
        assert_eq!(b.len(), b2.len());
    }
}