tests/test_rl_engine.py

# tests/test_rl_engine.py
"""Unit tests for the RL engine (Thompson Sampling bandit)."""

import json
import pytest
from unittest.mock import MagicMock, patch

from app.services.rl_engine import (
    ACTIONS,
    PARAM_CAP,
    PARAM_FLOOR,
    ACTION_FATIGUE_CEILING,
    _clamp_params,
    _get_default_params,
    _load_bandit_params,
    _compute_context_features,
    _context_adjusted_params,
    _thompson_sample,
    _apply_safety_guard,
    _detect_plateau,
    get_behavior_signal,
    update_reward,
)


# ──────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────

def make_user(
    user_id="testrl",
    mood=0.5, fatigue=0.0, sleep_quality=0.8,
    stress=0.1, shock=0.0, sick=0.0, competence=0.7,
    rl_bandit_params=None,
    rl_history_count=0,
    rl_state=None,
    current_streak=0,
    last_login_at=None,
):
    user = MagicMock()
    user.user_id          = user_id
    user.mood             = mood
    user.fatigue          = fatigue
    user.sleep_quality    = sleep_quality
    user.stress           = stress
    user.shock            = shock
    user.sick             = sick
    user.competence       = competence
    user.rl_bandit_params = rl_bandit_params
    user.rl_history_count = rl_history_count
    user.rl_state         = rl_state or []
    user.current_streak   = current_streak
    user.last_login_at    = last_login_at
    user.days_since_last_login = 0
    return user


# ──────────────────────────────────────────────
# Param loading
# ──────────────────────────────────────────────

def test_default_params_structure():
    params = _get_default_params()
    for action in ACTIONS:
        assert action in params
        assert "alpha" in params[action]
        assert "beta"  in params[action]


def test_load_bandit_params_fresh_user():
    user = make_user()
    params = _load_bandit_params(user)
    assert all(action in params for action in ACTIONS)


def test_load_bandit_params_valid_stored():
    stored = {a: {"alpha": 5.0, "beta": 3.0} for a in ACTIONS}
    user = make_user(rl_bandit_params=stored)
    params = _load_bandit_params(user)
    assert params["Recovery"]["alpha"] == 5.0


def test_load_bandit_params_corrupted_falls_back():
    user = make_user(rl_bandit_params="not_valid_json!!")
    params = _load_bandit_params(user)
    # Should not raise; should return defaults
    assert all(action in params for action in ACTIONS)


def test_clamp_params_enforces_bounds():
    bloated = {a: {"alpha": 999.0, "beta": 0.0001} for a in ACTIONS}
    clamped = _clamp_params(bloated)
    for action in ACTIONS:
        assert clamped[action]["alpha"] <= PARAM_CAP
        assert clamped[action]["beta"]  >= PARAM_FLOOR


# ──────────────────────────────────────────────
# Context features
# ──────────────────────────────────────────────

def test_context_features_have_expected_keys():
    user = make_user()
    state = [0.7, 0.5, 0.0, 0.8, 0.0, 0.0, 0.1]
    ctx = _compute_context_features(state, user, {})
    for key in ("energy", "disruption", "readiness", "time_factor", "weekend_factor",
                "streak_factor", "inactivity_factor", "competence", "mood", "fatigue"):
        assert key in ctx


def test_high_fatigue_context_boosts_disruption():
    user = make_user(fatigue=0.95, mood=0.3, sleep_quality=0.4)
    state = [0.7, 0.3, 0.95, 0.4, 0.0, 0.0, 0.1]
    ctx = _compute_context_features(state, user, {})
    assert ctx["energy"] < 0.5   # energy should be low with high fatigue, low mood, low sleep


# ──────────────────────────────────────────────
# Contextual adjustment
# ──────────────────────────────────────────────

def test_high_fatigue_boosts_recovery_arm():
    user = make_user(fatigue=0.9)
    base = _get_default_params()
    state = [0.7, 0.3, 0.9, 0.4, 0.0, 0.0, 0.8]
    ctx = _compute_context_features(state, user, {})
    _neutral_hist = {"consecutive_hard_days": 0, "recent_completion_rate": 0.7,
                     "avg_satisfaction_7d": 0.6, "fatigue_trend": 0.0,
                     "overwork_signal": 0.0, "deep_work_streak": 0.0, "avg_daily_minutes": 45.0}
    adjusted = _context_adjusted_params(base, ctx, _neutral_hist, is_cold_start=False, is_plateau=False)
    # Recovery should have a higher alpha than default
    assert adjusted["Recovery"]["alpha"] > base["Recovery"]["alpha"]


def test_plateau_boosts_exploration_arm():
    user = make_user()
    base = _get_default_params()
    state = [0.7, 0.5, 0.1, 0.8, 0.0, 0.0, 0.1]
    ctx = _compute_context_features(state, user, {})
    _neutral_hist = {"consecutive_hard_days": 0, "recent_completion_rate": 0.7,
                     "avg_satisfaction_7d": 0.6, "fatigue_trend": 0.0,
                     "overwork_signal": 0.0, "deep_work_streak": 0.0, "avg_daily_minutes": 45.0}
    adjusted_plateau    = _context_adjusted_params(base, ctx, _neutral_hist, is_cold_start=False, is_plateau=True)
    adjusted_no_plateau = _context_adjusted_params(base, ctx, _neutral_hist, is_cold_start=False, is_plateau=False)
    assert adjusted_plateau["Exploration"]["alpha"] > adjusted_no_plateau["Exploration"]["alpha"]


def test_cold_start_flattens_distribution():
    user = make_user()
    base = _get_default_params()
    state = [0.7, 0.5, 0.0, 0.8, 0.0, 0.0, 0.1]
    ctx = _compute_context_features(state, user, {})
    _neutral_hist = {"consecutive_hard_days": 0, "recent_completion_rate": 0.7,
                     "avg_satisfaction_7d": 0.6, "fatigue_trend": 0.0,
                     "overwork_signal": 0.0, "deep_work_streak": 0.0, "avg_daily_minutes": 45.0}
    adjusted_cold = _context_adjusted_params(base, ctx, _neutral_hist, is_cold_start=True,  is_plateau=False)
    adjusted_warm = _context_adjusted_params(base, ctx, _neutral_hist, is_cold_start=False, is_plateau=False)
    # Cold start should add beta to all arms (flatten the distribution)
    for action in ACTIONS:
        assert adjusted_cold[action]["beta"] >= adjusted_warm[action]["beta"]


# ──────────────────────────────────────────────
# Thompson Sampling
# ──────────────────────────────────────────────

def test_thompson_sample_returns_valid_action():
    params = {a: {"alpha": 2.0, "beta": 2.0} for a in ACTIONS}
    chosen, samples, propensity = _thompson_sample(params)
    assert chosen in ACTIONS
    assert set(samples.keys()) == set(ACTIONS)
    assert set(propensity.keys()) == set(ACTIONS)


def test_thompson_sample_biased_toward_high_alpha():
    """With very high alpha for Deep Work, should be chosen almost always."""
    params = {a: {"alpha": 1.0, "beta": 100.0} for a in ACTIONS}
    params["Deep Work"] = {"alpha": 100.0, "beta": 1.0}
    wins = sum(1 for _ in range(50) if _thompson_sample(params)[0] == "Deep Work")
    assert wins > 40, f"Expected Deep Work to win most rounds, got {wins}/50"


# ──────────────────────────────────────────────
# Safety guard
# ──────────────────────────────────────────────

def test_safety_guard_allows_recovery_at_any_fatigue():
    assert _apply_safety_guard("Recovery", 0.99) == "Recovery"


def test_safety_guard_downgrades_deep_work_at_high_fatigue():
    result = _apply_safety_guard("Deep Work", 0.9)
    assert result in ["Light Review", "Recovery"]
    assert result != "Deep Work"


def test_safety_guard_no_change_at_safe_fatigue():
    result = _apply_safety_guard("Deep Work", 0.3)   # Under 0.75 ceiling
    assert result == "Deep Work"


# ──────────────────────────────────────────────
# Plateau detection
# ──────────────────────────────────────────────

def test_plateau_detected_when_variance_low():
    user = make_user(rl_state=[0.7, 0.71, 0.69, 0.70, 0.71, 0.70, 0.69])
    assert _detect_plateau(user) is True


def test_no_plateau_when_variance_high():
    user = make_user(rl_state=[0.2, 0.9, 0.1, 0.8, 0.3, 0.95, 0.15])
    assert _detect_plateau(user) is False


def test_no_plateau_with_insufficient_history():
    user = make_user(rl_state=[0.5, 0.5, 0.5])  # Only 3 entries, need 7
    assert _detect_plateau(user) is False


def test_no_plateau_empty_history():
    user = make_user(rl_state=[])
    assert _detect_plateau(user) is False


# ──────────────────────────────────────────────
# get_behavior_signal (integration)
# ──────────────────────────────────────────────

def test_behavior_signal_is_crash_proof():
    """Should never raise regardless of user state."""
    user = make_user()
    user.rl_bandit_params = "corrupted{{{{json"
    signal = get_behavior_signal(user)
    assert "action_label" in signal
    assert signal["action_label"] in ACTIONS


def test_behavior_signal_structure():
    user = make_user()
    signal = get_behavior_signal(user)
    for key in ("action_label", "intensity", "mood", "fatigue",
                "rl_samples", "rl_updates", "cold_start", "plateau_detected"):
        assert key in signal


def test_behavior_signal_intensity_in_range():
    user = make_user()
    for _ in range(20):
        signal = get_behavior_signal(user)
        assert 0.0 <= signal["intensity"] <= 1.0


def test_cold_start_flag_on_new_user():
    user = make_user(rl_history_count=0)
    signal = get_behavior_signal(user)
    assert signal["cold_start"] is True


def test_no_cold_start_after_ten_updates():
    user = make_user(rl_history_count=15)
    signal = get_behavior_signal(user)
    assert signal["cold_start"] is False


# ──────────────────────────────────────────────
# update_reward
# ──────────────────────────────────────────────

def test_reward_update_increments_history_count():
    db = MagicMock()
    user = make_user(rl_history_count=5)
    user.rl_state = []
    update_reward(db, user, "Normal", 1.0)
    assert user.rl_history_count == 6


def test_reward_update_unknown_action_skipped():
    db = MagicMock()
    user = make_user()
    user.rl_state = []
    # Should not raise
    update_reward(db, user, "MadeUpAction", 0.5)
    assert user.rl_history_count == 0   # Not incremented


def test_reward_updates_alpha_on_success():
    db = MagicMock()
    user = make_user()
    user.rl_state = []
    update_reward(db, user, "Normal", 1.0)
    params = json.loads(user.rl_bandit_params)
    # Alpha for Normal should have increased
    assert params["Normal"]["alpha"] > _get_default_params()["Normal"]["alpha"] * 0.996


def test_reward_updates_beta_on_failure():
    db = MagicMock()
    user = make_user()
    user.rl_state = []
    update_reward(db, user, "Normal", 0.0)
    params = json.loads(user.rl_bandit_params)
    assert params["Normal"]["beta"] > _get_default_params()["Normal"]["beta"] * 0.996


def test_reward_clamped_to_valid_range():
    db = MagicMock()
    user = make_user()
    user.rl_state = []
    # Out-of-range reward should be clamped, not crash
    update_reward(db, user, "Normal", 5.0)
    params = json.loads(user.rl_bandit_params)
    assert params["Normal"]["alpha"] <= PARAM_CAP


def test_reward_history_appended_to_rl_state():
    db = MagicMock()
    user = make_user()
    user.rl_state = [0.5, 0.6]
    update_reward(db, user, "Normal", 0.8)
    assert len(user.rl_state) == 3
    assert user.rl_state[-1] == 0.8