# RECOMMENDED TEST ADDITIONS: Preventing Future Feature-Flag Bugs

## Quick Start Guide

### Step 1: Add These 3 Critical Tests IMMEDIATELY (2 hours)

These tests will expose the current bug and prevent similar issues:

```rust
// Add to ml/src/trainers/dqn.rs in the #[cfg(test)] mod tests block

/// CRITICAL TEST #1: Force epsilon > 0 to test exploration path
#[tokio::test]
async fn test_batched_action_selection_with_exploration() {
    let mut hyperparams = create_test_params();
    hyperparams.epsilon_start = 0.5;
    hyperparams.epsilon_end = 0.5;
    let trainer = DQNTrainer::new(hyperparams).unwrap();

    // Force 50% exploration
    {
        let mut agent = trainer.agent.write().await;
        agent.set_epsilon(0.5).unwrap();
    }

    let batch_size = 100;
    let mut states = Vec::with_capacity(batch_size);
    for i in 0..batch_size {
        let mut feature_vec = [0.0; 128];
        feature_vec[0] = 4000.0 + (i as f64 * 10.0);
        feature_vec[1] = 4010.0 + (i as f64 * 10.0);
        feature_vec[2] = 3990.0 + (i as f64 * 10.0);
        feature_vec[3] = 4005.0 + (i as f64 * 10.0);
        feature_vec[4] = 1000.0;
        for j in 5..128 { feature_vec[j] = (j as f64) * 0.1; }

        let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
            .unwrap_or(rust_decimal::Decimal::ZERO);
        let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price)).unwrap();
        states.push(state);
    }

    // THIS WILL FAIL with current bug (action_idx can be 0-44, TradingAction only accepts 0-2)
    let actions_result = trainer.select_actions_batch(&states).await;

    assert!(
        actions_result.is_ok(),
        "Exploration path should not crash: {:?}",
        actions_result.err()
    );

    let actions = actions_result.unwrap();
    assert_eq!(actions.len(), batch_size);
}

/// CRITICAL TEST #2: Force 100% exploration to maximize coverage
#[tokio::test]
async fn test_full_exploration_action_range() {
    let mut hyperparams = create_test_params();
    hyperparams.epsilon_start = 1.0;
    hyperparams.epsilon_end = 1.0;
    let trainer = DQNTrainer::new(hyperparams).unwrap();

    // Force 100% exploration
    {
        let mut agent = trainer.agent.write().await;
        agent.set_epsilon(1.0).unwrap();
    }

    let batch_size = 1000; // Large batch for statistical coverage
    let mut states = Vec::with_capacity(batch_size);
    for i in 0..batch_size {
        let mut feature_vec = [0.0; 128];
        feature_vec[0] = 4000.0 + (i as f64 * 10.0);
        feature_vec[1] = 4010.0 + (i as f64 * 10.0);
        feature_vec[2] = 3990.0 + (i as f64 * 10.0);
        feature_vec[3] = 4005.0 + (i as f64 * 10.0);
        feature_vec[4] = 1000.0;
        for j in 5..128 { feature_vec[j] = (j as f64) * 0.1; }

        let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
            .unwrap_or(rust_decimal::Decimal::ZERO);
        let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price)).unwrap();
        states.push(state);
    }

    // With 1000 samples and 100% exploration, will hit all action_idx values 0-44
    // THIS WILL DEFINITELY FAIL with current bug
    let actions_result = trainer.select_actions_batch(&states).await;

    assert!(
        actions_result.is_ok(),
        "Full exploration should not crash with any action_idx: {:?}",
        actions_result.err()
    );
}

/// CRITICAL TEST #3: Feature-specific validation
#[tokio::test]
#[cfg(feature = "factored-actions")]
async fn test_factored_actions_feature_validation() {
    let mut hyperparams = create_test_params();
    hyperparams.epsilon_start = 0.5;
    hyperparams.epsilon_end = 0.5;
    let trainer = DQNTrainer::new(hyperparams).unwrap();

    // Force exploration
    {
        let mut agent = trainer.agent.write().await;
        agent.set_epsilon(0.5).unwrap();
    }

    let batch_size = 500;
    let mut states = Vec::with_capacity(batch_size);
    for i in 0..batch_size {
        let mut feature_vec = [0.0; 128];
        feature_vec[0] = 4000.0 + (i as f64 * 10.0);
        feature_vec[1] = 4010.0 + (i as f64 * 10.0);
        feature_vec[2] = 3990.0 + (i as f64 * 10.0);
        feature_vec[3] = 4005.0 + (i as f64 * 10.0);
        feature_vec[4] = 1000.0;
        for j in 5..128 { feature_vec[j] = (j as f64) * 0.1; }

        let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
            .unwrap_or(rust_decimal::Decimal::ZERO);
        let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price)).unwrap();
        states.push(state);
    }

    let actions_result = trainer.select_actions_batch(&states).await;

    // This test is ONLY for factored-actions feature
    // If feature is enabled, we expect FactoredAction conversion, not TradingAction
    assert!(
        actions_result.is_ok(),
        "Factored-actions feature should handle 45-action space: {:?}",
        actions_result.err()
    );

    let actions = actions_result.unwrap();
    assert_eq!(actions.len(), batch_size);

    // TODO: Once refactored to return FactoredAction, add validation:
    // for action in actions.iter() {
    //     let factored = action.as_factored().expect("Should be FactoredAction");
    //     assert!(factored.exposure_level() <= 4);
    //     assert!(factored.order_type() <= 2);
    //     assert!(factored.urgency() <= 2);
    // }
}
```

### Step 2: Run Tests to Verify Bug Detection

```bash
# This should FAIL with current code (proves tests work)
cargo test --package ml --lib trainers::dqn::tests::test_batched_action_selection_with_exploration

# This should also FAIL
cargo test --package ml --lib trainers::dqn::tests::test_full_exploration_action_range

# This should also FAIL (factored-actions only)
cargo test --package ml --lib trainers::dqn::tests::test_factored_actions_feature_validation
```

**Expected output**:
```
thread 'trainers::dqn::tests::test_batched_action_selection_with_exploration' panicked at ml/src/trainers/dqn.rs:XXXX:
Exploration path should not crash: Some(Invalid action index: 29
```

### Step 3: Fix the Bug

After confirming tests detect the bug, fix the code (see FACTORED_ACTIONS_BUG_FIX_PLAN.md).

### Step 4: Verify Tests Pass After Fix

```bash
# All 3 new tests should pass
cargo test --package ml --lib trainers::dqn::tests -- --nocapture
```

---

## Test Matrix for CI/CD (1 hour)

Add to `.gitlab-ci.yml`:

```yaml
test-dqn-default:
  stage: test
  script:
    - cargo test -p ml --lib trainers::dqn::tests
  allow_failure: false

test-dqn-no-factored:
  stage: test
  script:
    - cargo test -p ml --lib trainers::dqn::tests --no-default-features --features cuda
  allow_failure: false

test-dqn-factored-explicit:
  stage: test
  script:
    - cargo test -p ml --lib trainers::dqn::tests --features cuda,factored-actions
  allow_failure: false
```

---

## Additional Recommended Tests (Medium Priority)

### Test 4: Action Diversity Validation (1 hour)

```rust
#[tokio::test]
#[cfg(not(feature = "factored-actions"))]
async fn test_exploration_action_diversity() {
    let mut hyperparams = create_test_params();
    hyperparams.epsilon_start = 1.0;
    hyperparams.epsilon_end = 1.0;
    let trainer = DQNTrainer::new(hyperparams).unwrap();

    {
        let mut agent = trainer.agent.write().await;
        agent.set_epsilon(1.0).unwrap();
    }

    let batch_size = 3000;
    let mut states = Vec::with_capacity(batch_size);
    for i in 0..batch_size {
        let mut feature_vec = [0.0; 128];
        feature_vec[0] = 4000.0 + (i as f64 * 10.0);
        feature_vec[1] = 4010.0 + (i as f64 * 10.0);
        feature_vec[2] = 3990.0 + (i as f64 * 10.0);
        feature_vec[3] = 4005.0 + (i as f64 * 10.0);
        feature_vec[4] = 1000.0;
        for j in 5..128 { feature_vec[j] = (j as f64) * 0.1; }

        let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
            .unwrap_or(rust_decimal::Decimal::ZERO);
        let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price)).unwrap();
        states.push(state);
    }

    let actions = trainer.select_actions_batch(&states).await.unwrap();

    // Count action distribution
    let mut buy_count = 0;
    let mut sell_count = 0;
    let mut hold_count = 0;
    for action in actions.iter() {
        match action {
            TradingAction::Buy => buy_count += 1,
            TradingAction::Sell => sell_count += 1,
            TradingAction::Hold => hold_count += 1,
        }
    }

    // With 3000 samples and 3 actions, expect ~1000 of each
    // Use 20% tolerance (800-1200 range)
    let expected = 1000;
    let tolerance = 200;

    assert!(
        buy_count >= expected - tolerance && buy_count <= expected + tolerance,
        "Expected ~{} Buy actions, got {} (distribution: Buy={}, Sell={}, Hold={})",
        expected, buy_count, buy_count, sell_count, hold_count
    );

    assert!(
        sell_count >= expected - tolerance && sell_count <= expected + tolerance,
        "Expected ~{} Sell actions, got {} (distribution: Buy={}, Sell={}, Hold={})",
        expected, sell_count, buy_count, sell_count, hold_count
    );

    assert!(
        hold_count >= expected - tolerance && hold_count <= expected + tolerance,
        "Expected ~{} Hold actions, got {} (distribution: Buy={}, Sell={}, Hold={})",
        expected, hold_count, buy_count, sell_count, hold_count
    );

    println!("✅ Action diversity validated: Buy={}, Sell={}, Hold={}",
             buy_count, sell_count, hold_count);
}
```

### Test 5: Epsilon Boundary Conditions (30 min)

```rust
#[tokio::test]
async fn test_epsilon_boundary_conditions() {
    // Test epsilon=0.0 (pure exploitation)
    let mut hyperparams = create_test_params();
    hyperparams.epsilon_start = 0.0;
    hyperparams.epsilon_end = 0.0;
    let trainer = DQNTrainer::new(hyperparams).unwrap();

    {
        let mut agent = trainer.agent.write().await;
        agent.set_epsilon(0.0).unwrap();
    }

    let batch_size = 10;
    let mut states = Vec::with_capacity(batch_size);
    for i in 0..batch_size {
        let mut feature_vec = [0.0; 128];
        feature_vec[0] = 4000.0 + (i as f64 * 10.0);
        feature_vec[1] = 4010.0 + (i as f64 * 10.0);
        feature_vec[2] = 3990.0 + (i as f64 * 10.0);
        feature_vec[3] = 4005.0 + (i as f64 * 10.0);
        feature_vec[4] = 1000.0;
        for j in 5..128 { feature_vec[j] = (j as f64) * 0.1; }

        let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
            .unwrap_or(rust_decimal::Decimal::ZERO);
        let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price)).unwrap();
        states.push(state);
    }

    let actions_0 = trainer.select_actions_batch(&states).await.unwrap();
    assert_eq!(actions_0.len(), batch_size);

    // Test epsilon=1.0 (pure exploration)
    {
        let mut agent = trainer.agent.write().await;
        agent.set_epsilon(1.0).unwrap();
    }

    let actions_1 = trainer.select_actions_batch(&states).await.unwrap();
    assert_eq!(actions_1.len(), batch_size);

    println!("✅ Epsilon boundary conditions validated (0.0 and 1.0)");
}
```

---

## Property-Based Testing (Optional, 3-4 hours)

For maximum robustness, add property-based tests:

```rust
// Add to ml/Cargo.toml [dev-dependencies]
proptest = "1.5"

// Add to ml/src/trainers/dqn.rs
#[cfg(test)]
mod proptests {
    use super::*;
    use proptest::prelude::*;

    proptest! {
        #[test]
        fn prop_action_selection_never_panics(
            epsilon in 0.0..=1.0f64,
            batch_size in 1..=100usize
        ) {
            tokio::runtime::Runtime::new().unwrap().block_on(async {
                let mut hyperparams = create_test_params();
                hyperparams.epsilon_start = epsilon;
                hyperparams.epsilon_end = epsilon;
                let trainer = DQNTrainer::new(hyperparams).unwrap();

                {
                    let mut agent = trainer.agent.write().await;
                    agent.set_epsilon(epsilon).unwrap();
                }

                let mut states = Vec::with_capacity(batch_size);
                for i in 0..batch_size {
                    let mut feature_vec = [0.0; 128];
                    feature_vec[0] = 4000.0 + (i as f64 * 10.0);
                    feature_vec[1] = 4010.0 + (i as f64 * 10.0);
                    feature_vec[2] = 3990.0 + (i as f64 * 10.0);
                    feature_vec[3] = 4005.0 + (i as f64 * 10.0);
                    feature_vec[4] = 1000.0;
                    for j in 5..128 { feature_vec[j] = (j as f64) * 0.1; }

                    let close_price = rust_decimal::Decimal::try_from(feature_vec[3])
                        .unwrap_or(rust_decimal::Decimal::ZERO);
                    let state = trainer.feature_vector_to_state(&feature_vec, Some(close_price)).unwrap();
                    states.push(state);
                }

                // Property: Action selection should never panic for any epsilon or batch_size
                let result = trainer.select_actions_batch(&states).await;
                prop_assert!(result.is_ok(), "Action selection panicked: {:?}", result.err());

                let actions = result.unwrap();
                prop_assert_eq!(actions.len(), batch_size);
            });
        }
    }
}
```

---

## Summary of Effort

| Priority | Tests | Effort | Benefit |
|---------|-------|--------|---------|
| **CRITICAL (Add NOW)** | 3 tests | 2 hours | Catches current bug + prevents recurrence |
| **HIGH (Add this week)** | CI/CD matrix | 1 hour | Prevents feature-flag bugs |
| **MEDIUM (Add this sprint)** | 2 tests | 1.5 hours | Improves coverage |
| **OPTIONAL (Add if time)** | Property tests | 3-4 hours | Maximum robustness |
| **TOTAL** | 5-6 tests + CI | 4.5-7.5 hours | Production-ready test suite |

---

## Validation Checklist

After adding tests, verify:

- [ ] Tests FAIL with current code (proves detection works)
- [ ] Tests PASS after bug fix (proves fix works)
- [ ] Tests run in CI/CD pipeline (prevents regression)
- [ ] Tests cover both feature flags (factored-actions ON/OFF)
- [ ] Tests cover epsilon boundaries (0.0, 0.5, 1.0)
- [ ] Test output includes clear error messages
- [ ] Tests are documented with comments explaining purpose

---

## Next Steps

1. **Copy tests from this document to `ml/src/trainers/dqn.rs`**
2. **Run tests to confirm they FAIL** (proves detection)
3. **Fix the bug** (see FACTORED_ACTIONS_BUG_FIX_PLAN.md)
4. **Re-run tests to confirm they PASS** (proves fix)
5. **Add CI/CD test matrix** (prevents regression)
6. **Optional: Add property-based tests** (maximum robustness)

**Total time investment**: 4.5-7.5 hours
**Benefit**: Never ship a feature-flag bug to production again