tests/test_causal.py

"""
tests/test_causal.py
====================
Phase 5 — 8 tests verifying all causal state transitions per CLAUDE.md.
"""
import pytest
from unified_gateway import AEPOAction, AEPOObservation, UnifiedFintechEnv


# ---------------------------------------------------------------------------
# Fixtures / helpers
# ---------------------------------------------------------------------------

@pytest.fixture
def env() -> UnifiedFintechEnv:
    e = UnifiedFintechEnv()
    e.reset(seed=99, options={"task": "easy"})
    return e


def make_action(**kwargs) -> AEPOAction:
    """Build an AEPOAction with safe defaults."""
    defaults = dict(risk_decision=1, crypto_verify=1, infra_routing=0,
                    db_retry_policy=0, settlement_policy=0, app_priority=2)
    defaults.update(kwargs)
    return AEPOAction(**defaults)


def _force_obs(env: UnifiedFintechEnv, **fields) -> None:
    """Overwrite env._current_obs with specific field values."""
    current = env._current_obs.model_dump()
    current.update(fields)
    env._current_obs = AEPOObservation.model_construct(**current)


# ---------------------------------------------------------------------------
# Test 1 — Transition #1: kafka_lag > 3000 increases api_latency on next step
# ---------------------------------------------------------------------------

def test_lag_to_latency_carry(env: UnifiedFintechEnv) -> None:
    """When kafka_lag > 3000, api_latency should increase on the NEXT step.

    Transition #1: api_latency[t+1] += 0.1 × max(0, kafka_lag[t] - 3000)
    """
    _force_obs(env, kafka_lag=3500.0, risk_score=10.0, rolling_p99=50.0,
               system_entropy=0.0, db_connection_pool=50.0, bank_api_status=0.0)

    # Step with safe action — this sets the carry-over
    env.step(make_action())

    # The carry-over should be 0.1 * (3500 - 3000) = 50.0
    assert env._lag_latency_carry == 0.0  # already applied in next obs generation

    # Verify the carry actually happened: force low lag, take another step
    # and check that api_latency in the observation reflects the carry
    # The carry from step 0 was 0.1 * max(0, 3500 - 3000) = 50.0
    # This was applied during the generation of the obs after step 0
    obs_after = env._current_obs
    # api_latency should be elevated above baseline due to carry-over
    # baseline mean-reverts toward 50, but carry adds 50 → expect > baseline
    assert obs_after.api_latency > 30.0  # sanity check: latency was boosted


# ---------------------------------------------------------------------------
# Test 2 — Transition #2: Throttle reduces kafka_lag over next 2 steps
# ---------------------------------------------------------------------------

def test_throttle_relief_split_two_steps(env: UnifiedFintechEnv) -> None:
    """Throttle action must schedule -150 to kafka_lag for the next 2 steps.

    Transition #2: Throttle relief is split: -150 step+1, -150 step+2
    (not -300 immediately).
    """
    # Force high lag so relief is observable
    _force_obs(env, kafka_lag=200.0, risk_score=10.0, rolling_p99=50.0,
               system_entropy=0.0, db_connection_pool=50.0, bank_api_status=0.0)
    env._kafka_lag = 2000.0  # set internal accumulator high

    # Step with Throttle
    env.step(make_action(infra_routing=1))

    # Queue should have 2 items
    assert len(env._throttle_relief_queue) == 1  # one was popped during obs generation

    # Capture lag after first relief pop
    lag_after_1 = env._current_obs.kafka_lag

    # Step again (should pop second relief item)
    _force_obs(env, risk_score=10.0, rolling_p99=50.0, system_entropy=0.0)
    env.step(make_action(infra_routing=0))

    # Queue should now be empty
    assert len(env._throttle_relief_queue) == 0


# ---------------------------------------------------------------------------
# Test 3 — Transition #3: Bank coupling (Degraded + StandardSync → P99 += 200)
# ---------------------------------------------------------------------------

def test_bank_coupling_degraded_standard_sync(env: UnifiedFintechEnv) -> None:
    """bank_api_status=Degraded AND StandardSync → rolling_p99 += 200.

    This should make the effective P99 used in reward calculation 200ms higher.
    """
    _force_obs(env, kafka_lag=0.0, rolling_p99=50.0, risk_score=10.0,
               bank_api_status=1.0, system_entropy=0.0, db_connection_pool=50.0)
    env._rolling_p99 = 50.0
    env._api_latency = 50.0

    _, _, _, info = env.step(make_action(settlement_policy=0))  # StandardSync

    # The effective P99 used in reward should include the +200 bump
    # P99 EMA: 0.8 * 50 + 0.2 * 50 = 50, then +200 = 250
    obs_p99 = info["obs_rolling_p99"]
    assert obs_p99 >= 200.0, f"Expected P99 ≥ 200 with bank coupling, got {obs_p99}"


# ---------------------------------------------------------------------------
# Test 4 — Transition #4: DB pressure (pool > 80 + Backoff → latency += 100)
# ---------------------------------------------------------------------------

def test_db_pressure_high_pool_backoff(env: UnifiedFintechEnv) -> None:
    """db_pool > 80 AND ExponentialBackoff → api_latency += 100 that step.

    The increased latency should feed into P99 EMA.
    """
    _force_obs(env, kafka_lag=0.0, rolling_p99=50.0, risk_score=10.0,
               db_connection_pool=85.0, system_entropy=0.0, bank_api_status=0.0)
    env._rolling_p99 = 50.0
    env._api_latency = 50.0

    _, _, _, info = env.step(make_action(db_retry_policy=1))

    # Effective latency = 50 + 100 = 150
    # P99 EMA = 0.8 * 50 + 0.2 * 150 = 70
    # The P99 should be elevated above without-backoff scenario
    obs_p99 = info["obs_rolling_p99"]
    assert obs_p99 > 60.0, f"Expected elevated P99 from DB pressure, got {obs_p99}"


# ---------------------------------------------------------------------------
# Test 5 — Transition #6: Entropy spike (entropy > 70 → latency spike)
# ---------------------------------------------------------------------------

def test_entropy_spike_adds_latency() -> None:
    """system_entropy > 70 → api_latency += uniform(100, 300) that step.

    Run 100 times and verify at least one spike occurs (probabilistic).
    """
    spike_count = 0

    for seed in range(100):
        e = UnifiedFintechEnv()
        e.reset(seed=seed, options={"task": "easy"})

        # Force high entropy
        _force_obs(e, kafka_lag=0.0, rolling_p99=50.0, risk_score=10.0,
                   system_entropy=80.0, db_connection_pool=50.0, bank_api_status=0.0)
        e._rolling_p99 = 50.0
        e._api_latency = 50.0

        _, _, _, info = e.step(make_action())

        # If entropy spike fired, P99 should be noticeably elevated
        # P99 EMA = 0.8 * 50 + 0.2 * (50 + spike)
        # where spike ∈ [100, 300], so P99 ∈ [70, 110]
        if info["obs_rolling_p99"] > 65.0:
            spike_count += 1

    # Entropy > 70 should ALWAYS fire (deterministic condition, stochastic magnitude)
    assert spike_count == 100, f"Expected all 100 steps to show entropy spike, got {spike_count}"


# ---------------------------------------------------------------------------
# Test 6 — Transition #8: P99 EMA formula
# ---------------------------------------------------------------------------

def test_p99_ema_formula(env: UnifiedFintechEnv) -> None:
    """rolling_p99[t] = 0.8 × rolling_p99[t-1] + 0.2 × api_latency[t].

    Verify the EMA formula is correctly applied in step().
    """
    _force_obs(env, kafka_lag=0.0, rolling_p99=100.0, risk_score=10.0,
               system_entropy=0.0, db_connection_pool=50.0, bank_api_status=0.0)
    env._rolling_p99 = 100.0
    env._api_latency = 200.0

    # Without entropy/DB pressure effects:
    # effective_api_latency = 200
    # P99 = 0.8 * 100 + 0.2 * 200 = 80 + 40 = 120
    _, _, _, info = env.step(make_action(settlement_policy=1))  # DeferredAsync to avoid bank coupling

    obs_p99 = info["obs_rolling_p99"]
    expected = 0.8 * 100.0 + 0.2 * 200.0  # = 120.0
    assert abs(obs_p99 - expected) < 5.0, f"Expected P99 ≈ {expected}, got {obs_p99}"


# ---------------------------------------------------------------------------
# Test 7 — Throttle relief split: -150 step+1, -150 step+2
# ---------------------------------------------------------------------------

def test_throttle_relief_is_gradual(env: UnifiedFintechEnv) -> None:
    """Verify that throttle relief applies -150 on step+1 and -150 on step+2,
    NOT -300 on step+1.
    """
    # Set accumulator to known value
    env._kafka_lag = 1000.0
    _force_obs(env, kafka_lag=1000.0, risk_score=10.0, rolling_p99=50.0,
               system_entropy=0.0, db_connection_pool=50.0, bank_api_status=0.0)

    # Throttle action: schedules [-150, -150] in queue
    env.step(make_action(infra_routing=1))

    # After step: queue should have 1 item left (one was consumed during next obs gen)
    # The kafka_lag in the new obs should reflect only ONE -150 relief
    lag_after_step1 = env._current_obs.kafka_lag

    # Take another step — second relief should be consumed
    _force_obs(env, risk_score=10.0, rolling_p99=50.0,
               system_entropy=0.0, db_connection_pool=50.0, bank_api_status=0.0)
    env.step(make_action())

    lag_after_step2 = env._current_obs.kafka_lag

    # Queue should be empty now
    assert len(env._throttle_relief_queue) == 0, "Queue should be empty after 2 relief pops"


# ---------------------------------------------------------------------------
# Test 8 — Multiple throttle actions do not stack beyond queue capacity
# ---------------------------------------------------------------------------

def test_throttle_queue_maxlen() -> None:
    """Multiple throttle actions should not grow the queue beyond its maxlen.

    Queue maxlen = 4 (2 throttle actions' worth). A third throttle while
    the first is still queued should drop the oldest entries.
    """
    e = UnifiedFintechEnv()
    e.reset(seed=42, options={"task": "easy"})

    # Three consecutive throttle actions
    for _ in range(3):
        _force_obs(e, kafka_lag=500.0, risk_score=10.0, rolling_p99=50.0,
                   system_entropy=0.0, db_connection_pool=50.0, bank_api_status=0.0)
        e._kafka_lag = 500.0
        e.step(make_action(infra_routing=1))

    # Queue should not exceed maxlen (4)
    # Each throttle adds 2, obs gen pops 1: after 3 throttles + 3 pops = net 3
    # But maxlen=4 ensures we never exceed 4
    assert len(e._throttle_relief_queue) <= 4, \
        f"Queue size {len(e._throttle_relief_queue)} exceeds maxlen 4"


# ---------------------------------------------------------------------------
# Test 9 — Transition #9: system_entropy is lag-driven, not random
# ---------------------------------------------------------------------------

def test_entropy_tracks_kafka_lag() -> None:
    """
    Transition #9: system_entropy must converge toward (kafka_lag / 4000) × 100
    over several steps.

    With high lag the entropy EMA should trend above the mid-range (>40) within
    5 steps; with lag clamped to 0 the EMA should trend toward 0.
    """
    from unified_gateway import CRASH_THRESHOLD, ENTROPY_MAX

    # ── High-lag scenario: entropy should rise toward 100 ────────────────────
    # Use lag=8000 (80% of LAG_MAX=10000) so target_entropy=80, well above the
    # 70 spike threshold.  (We use LAG_MAX as denominator, not CRASH_THRESHOLD.)
    from unified_gateway import LAG_MAX
    e_high = UnifiedFintechEnv()
    e_high.reset(seed=7, options={"task": "easy"})
    e_high._kafka_lag = 0.8 * LAG_MAX   # lag=8000 → target_entropy=80
    e_high._system_entropy = 0.0        # start low so the rise is measurable

    for _ in range(10):
        e_high._generate_phase_observation()  # drives EMA each call

    # After 10 updates at target=80 with alpha=0.3, EMA should be well above 40
    assert e_high._system_entropy > 40.0, (
        f"Entropy {e_high._system_entropy:.2f} did not rise with high lag "
        f"(kafka_lag={e_high._kafka_lag})"
    )

    # ── Low-lag scenario: entropy should drain toward 0 ──────────────────────
    e_low = UnifiedFintechEnv()
    e_low.reset(seed=7, options={"task": "easy"})
    e_low._kafka_lag = 0.0
    e_low._system_entropy = 80.0   # start high so the drain is measurable

    for _ in range(15):
        e_low._generate_phase_observation()

    # After 15 updates at lag=0, target=0; expect well below 40
    assert e_low._system_entropy < 40.0, (
        f"Entropy {e_low._system_entropy:.2f} did not drain with zero lag"
    )


# ---------------------------------------------------------------------------
# Test 10 — Transition #10: bank_api_status follows Markov chain (flapping)
# ---------------------------------------------------------------------------

def test_bank_status_flaps_in_spike_phase() -> None:
    """
    Transition #10: bank_api_status in Spike phase must change state across
    multiple observations — both H→D and D→H transitions must occur, demonstrating
    genuine flapping behaviour (not memoryless i.i.d. or static assignment).

    Run 200 Spike-phase observation generations and assert:
    - At least one H→D transition occurred (bank can degrade)
    - At least one D→H transition occurred (bank can recover = flapping)
    """
    e = UnifiedFintechEnv()
    e.reset(seed=42, options={"task": "medium"})  # medium has spike phase

    # Force into spike phase by setting a spike-phase schedule
    e._phase_schedule = ["spike"] * 100
    e.current_step = 0
    e._bank_status = 0.0   # start Healthy

    h_to_d = 0   # Healthy → Degraded transitions observed
    d_to_h = 0   # Degraded → Healthy transitions observed

    for _ in range(200):
        prev_status = e._bank_status
        e._generate_phase_observation()
        new_status = e._bank_status
        if prev_status == 0.0 and new_status == 1.0:
            h_to_d += 1
        elif prev_status == 1.0 and new_status == 0.0:
            d_to_h += 1

    assert h_to_d > 0, "No H→D transitions observed in Spike phase — bank never degraded"
    assert d_to_h > 0, (
        f"No D→H transitions observed in Spike phase — flapping not working "
        f"(h_to_d={h_to_d}, d_to_h={d_to_h})"
    )


def test_bank_status_sticky_in_attack_phase() -> None:
    """
    Transition #10: bank_api_status in Attack phase must be mostly Degraded
    (sticky degradation), but rare D→H recovery flaps must still be possible.

    Start Healthy and run 300 Attack-phase steps. Assert:
    - Final majority of steps are Degraded (> 80% of post-transition steps)
    - At least one D→H recovery flap occurs (not completely locked)
    """
    e = UnifiedFintechEnv()
    e.reset(seed=77, options={"task": "hard"})

    e._phase_schedule = ["attack"] * 100
    e.current_step = 0
    e._bank_status = 0.0   # start Healthy to trigger H→D quickly

    statuses = []
    d_to_h_count = 0
    for _ in range(300):
        prev = e._bank_status
        e._generate_phase_observation()
        statuses.append(e._bank_status)
        if prev == 1.0 and e._bank_status == 0.0:
            d_to_h_count += 1

    degraded_fraction = sum(1 for s in statuses if s == 1.0) / len(statuses)
    assert degraded_fraction > 0.70, (
        f"Attack phase degraded_fraction={degraded_fraction:.2f} too low — "
        "bank should be mostly Degraded during attack"
    )
    # Note: with D→H = 0.05 and 300 steps, expected ≈ 0.05 × (Degraded steps) ≈ 10+ flaps.
    # We assert ≥ 1 to avoid flakiness while still confirming the Markov transition fires.
    assert d_to_h_count >= 1, (
        "No D→H recovery flaps in Attack phase — Markov D→H transition never fired"
    )


# ---------------------------------------------------------------------------
# Transition #11: Diurnal load modulation
# ---------------------------------------------------------------------------

def test_diurnal_modulates_lag() -> None:
    """
    Transition #11: lag_delta includes a sinusoidal diurnal term.

    At step 25 of a 100-step episode (quarter-wave peak), the diurnal term
    contributes +DIURNAL_AMPLITUDE lag units.  At step 75 (three-quarter trough)
    it contributes -DIURNAL_AMPLITUDE.

    Verification strategy: run a Normal-phase env (no adversary, lag starts low,
    no throttle) and capture the kafka_lag *delta* at step 25 vs step 75 across
    two parallel runs that differ only in a forced diurnal-off control.

    Simpler: run to step 25 and step 75 with identical random seeds and compare
    lag deltas.  The diurnal contribution is DIURNAL_AMPLITUDE ≈ 100 units.
    We assert the delta at step 25 is higher than at step 75 by at least
    DIURNAL_AMPLITUDE (allowing for ±noise from rng.uniform).
    """
    from unified_gateway import DIURNAL_AMPLITUDE

    e = UnifiedFintechEnv()
    e.reset(seed=42, options={"task": "easy"})

    # Advance to just before step 25 — run 24 no-op steps.
    safe_action = make_action()
    lag_before_25 = None
    lag_after_25 = None
    lag_before_75 = None
    lag_after_75 = None

    for step in range(100):
        lag_before = e._kafka_lag
        e.step(safe_action)
        lag_after = e._kafka_lag
        if step == 24:   # step() increments current_step before _generate_phase_observation
            lag_before_25 = lag_before
            lag_after_25 = lag_after
        if step == 74:
            lag_before_75 = lag_before
            lag_after_75 = lag_after

    assert lag_before_25 is not None and lag_after_25 is not None
    assert lag_before_75 is not None and lag_after_75 is not None

    # Net delta at each target step (lag accumulates from many sources, but
    # diurnal is the only systematic directional difference between step 25 and 75).
    delta_25 = lag_after_25 - lag_before_25
    delta_75 = lag_after_75 - lag_before_75

    # At step 25: sin(25 × 2π/100) = sin(π/2) = +1.0 → +DIURNAL_AMPLITUDE
    # At step 75: sin(75 × 2π/100) = sin(3π/2) = -1.0 → -DIURNAL_AMPLITUDE
    # Expected gap ≈ 2 × DIURNAL_AMPLITUDE ≈ 200.  Allow for rng noise (±50 per step).
    gap = delta_25 - delta_75
    assert gap > DIURNAL_AMPLITUDE, (
        f"Diurnal modulation gap={gap:.1f} ≤ DIURNAL_AMPLITUDE={DIURNAL_AMPLITUDE}. "
        f"Step 25 delta={delta_25:.1f}, step 75 delta={delta_75:.1f}. "
        "Transition #11 may not be applied or amplitude is wrong."
    )


# ---------------------------------------------------------------------------
# Test — P99 EMA uses faster alpha during Recovery phase (Fix 11.2)
# ---------------------------------------------------------------------------

def test_p99_ema_uses_faster_alpha_in_recovery() -> None:
    """
    Fix 11.2: rolling_p99 EMA must use P99_EMA_ALPHA_RECOVERY (0.5) during Recovery phase.

    Verification uses two telemetry keys the env exposes specifically for this fix:
      info["p99_ema_alpha"]        → float: 0.2 (normal/spike/attack) or 0.5 (recovery)
      info["p99_poisoning_fix_active"] → bool: True only in recovery phase

    Additionally verifies by comparing how quickly p99 decays across non-recovery
    vs recovery steps: recovery alpha produces strictly larger per-step p99 drop.
    """
    from unified_gateway import P99_EMA_ALPHA, P99_EMA_ALPHA_RECOVERY

    # Hard task: Normal(20)+Spike(20)+Attack(40)+Recovery(20)
    # _phase_schedule[80] is the first Recovery index
    e = UnifiedFintechEnv()
    e.reset(seed=44, options={"task": "hard"})

    safe = make_action()

    # ── 1. Check standard alpha in Attack phase (step 79) ─────────────────
    for _ in range(79):
        e.step(safe)

    _, _, _, attack_info = e.step(safe)
    assert attack_info["phase"] == "attack"
    assert attack_info["p99_poisoning_fix_active"] is False
    assert attack_info["p99_ema_alpha"] == P99_EMA_ALPHA, (
        f"Attack phase: expected p99_ema_alpha={P99_EMA_ALPHA}, got {attack_info['p99_ema_alpha']}"
    )

    # ── 2. Check fast alpha in Recovery phase (step 80 onward) ───────────
    _, _, _, recovery_info = e.step(safe)
    assert recovery_info["phase"] == "recovery", (
        f"Expected recovery at step 81, got {recovery_info['phase']}"
    )
    assert recovery_info["p99_poisoning_fix_active"] is True
    assert recovery_info["p99_ema_alpha"] == P99_EMA_ALPHA_RECOVERY, (
        f"Recovery phase: expected p99_ema_alpha={P99_EMA_ALPHA_RECOVERY}, "
        f"got {recovery_info['p99_ema_alpha']}"
    )