//! Tests for batch loading routine concurrent counts (N+1 query fix). //! //! Verifies: //! 1. Batch query returns correct counts for multiple routines //! 2. Concurrent limit enforcement uses batch counts correctly #[cfg(feature = "libsql")] mod tests { use std::sync::Arc; use chrono::Utc; use uuid::Uuid; use ironclaw::agent::routine::{ Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger, }; use ironclaw::db::Database; async fn create_test_db() -> (Arc, tempfile::TempDir) { use ironclaw::db::libsql::LibSqlBackend; let temp_dir = tempfile::tempdir().expect("tempdir"); let db_path = temp_dir.path().join("test.db"); let backend = LibSqlBackend::new_local(&db_path) .await .expect("LibSqlBackend"); backend.run_migrations().await.expect("migrations"); let db: Arc = Arc::new(backend); (db, temp_dir) } // ----------------------------------------------------------------------- // Test 1: Batch query returns correct counts for multiple routines // ----------------------------------------------------------------------- #[tokio::test] async fn batch_query_empty_list() { let (db, _tmp) = create_test_db().await; let counts = db .count_running_routine_runs_batch(&[]) .await .expect("batch query should not fail"); assert!(counts.is_empty(), "Empty input should return empty map"); } #[tokio::test] async fn batch_query_single_routine() { let (db, _tmp) = create_test_db().await; let routine_id = Uuid::new_v4(); // Create routine let routine = Routine { id: routine_id, name: "test-routine".to_string(), description: "Test".to_string(), user_id: "default".to_string(), enabled: true, trigger: Trigger::Cron { schedule: "* * * * *".to_string(), timezone: None, }, action: RoutineAction::Lightweight { prompt: "test".to_string(), context_paths: vec![], max_tokens: 1000, use_tools: false, max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(0), max_concurrent: 5, dedup_window: None, }, notify: Default::default(), last_run_at: None, next_fire_at: None, run_count: 0, consecutive_failures: 0, state: serde_json::json!({}), created_at: Utc::now(), updated_at: Utc::now(), }; db.create_routine(&routine).await.expect("create routine"); // Create 3 running runs for _ in 0..3 { let run = RoutineRun { id: Uuid::new_v4(), routine_id, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: None, status: RunStatus::Running, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); } // Batch query for single routine let counts = db .count_running_routine_runs_batch(&[routine_id]) .await .expect("batch query should work"); assert_eq!(counts.len(), 1, "Should return 1 routine"); assert_eq!(counts[&routine_id], 3, "Should count 3 running runs"); } #[tokio::test] async fn batch_query_multiple_routines_different_counts() { let (db, _tmp) = create_test_db().await; let r1 = Uuid::new_v4(); let r2 = Uuid::new_v4(); let r3 = Uuid::new_v4(); // Create 3 routines for routine_id in [r1, r2, r3] { let routine = Routine { id: routine_id, name: format!("routine-{}", routine_id), description: "Test".to_string(), user_id: "default".to_string(), enabled: true, trigger: Trigger::Cron { schedule: "* * * * *".to_string(), timezone: None, }, action: RoutineAction::Lightweight { prompt: "test".to_string(), context_paths: vec![], max_tokens: 1000, use_tools: false, max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(0), max_concurrent: 5, dedup_window: None, }, notify: Default::default(), last_run_at: None, next_fire_at: None, run_count: 0, consecutive_failures: 0, state: serde_json::json!({}), created_at: Utc::now(), updated_at: Utc::now(), }; db.create_routine(&routine).await.expect("create routine"); } // r1: 2 running for _ in 0..2 { let run = RoutineRun { id: Uuid::new_v4(), routine_id: r1, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: None, status: RunStatus::Running, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); } // r2: 1 running let run = RoutineRun { id: Uuid::new_v4(), routine_id: r2, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: None, status: RunStatus::Running, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); // r3: 0 running (but has 1 Ok result) let run = RoutineRun { id: Uuid::new_v4(), routine_id: r3, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: Some(Utc::now()), status: RunStatus::Ok, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); // Single batch query for all 3 let counts = db .count_running_routine_runs_batch(&[r1, r2, r3]) .await .expect("batch query should work"); assert_eq!(counts.len(), 3, "Should return 3 routines"); assert_eq!(counts[&r1], 2, "r1 should have 2 running"); assert_eq!(counts[&r2], 1, "r2 should have 1 running"); assert_eq!( counts[&r3], 0, "r3 should have 0 running (Ok status is not running)" ); } #[tokio::test] async fn batch_query_missing_routines_default_to_zero() { let (db, _tmp) = create_test_db().await; let r1 = Uuid::new_v4(); let r2 = Uuid::new_v4(); let r3 = Uuid::new_v4(); // This one won't exist // Only create r1 let routine = Routine { id: r1, name: "routine-1".to_string(), description: "Test".to_string(), user_id: "default".to_string(), enabled: true, trigger: Trigger::Cron { schedule: "* * * * *".to_string(), timezone: None, }, action: RoutineAction::Lightweight { prompt: "test".to_string(), context_paths: vec![], max_tokens: 1000, use_tools: false, max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(0), max_concurrent: 5, dedup_window: None, }, notify: Default::default(), last_run_at: None, next_fire_at: None, run_count: 0, consecutive_failures: 0, state: serde_json::json!({}), created_at: Utc::now(), updated_at: Utc::now(), }; db.create_routine(&routine).await.expect("create routine"); // r1 has 1 running let run = RoutineRun { id: Uuid::new_v4(), routine_id: r1, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: None, status: RunStatus::Running, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); // Query for r1, r2 (doesn't exist), r3 (doesn't exist) let counts = db .count_running_routine_runs_batch(&[r1, r2, r3]) .await .expect("batch query should work"); assert_eq!(counts.len(), 3, "Should have all 3 routine IDs"); assert_eq!(counts[&r1], 1, "r1 should have 1 running"); assert_eq!(counts[&r2], 0, "r2 should default to 0"); assert_eq!(counts[&r3], 0, "r3 should default to 0"); } #[tokio::test] async fn batch_query_only_counts_running_status() { let (db, _tmp) = create_test_db().await; let routine_id = Uuid::new_v4(); // Create routine let routine = Routine { id: routine_id, name: "test-routine".to_string(), description: "Test".to_string(), user_id: "default".to_string(), enabled: true, trigger: Trigger::Cron { schedule: "* * * * *".to_string(), timezone: None, }, action: RoutineAction::Lightweight { prompt: "test".to_string(), context_paths: vec![], max_tokens: 1000, use_tools: false, max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(0), max_concurrent: 5, dedup_window: None, }, notify: Default::default(), last_run_at: None, next_fire_at: None, run_count: 0, consecutive_failures: 0, state: serde_json::json!({}), created_at: Utc::now(), updated_at: Utc::now(), }; db.create_routine(&routine).await.expect("create routine"); // Create 5 runs with mixed statuses let statuses = [ RunStatus::Running, RunStatus::Running, RunStatus::Ok, RunStatus::Failed, RunStatus::Attention, ]; for status in statuses.iter() { let run = RoutineRun { id: Uuid::new_v4(), routine_id, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: Some(Utc::now()), status: *status, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); } // Batch query should only count Running status let counts = db .count_running_routine_runs_batch(&[routine_id]) .await .expect("batch query should work"); assert_eq!( counts[&routine_id], 2, "Should only count 2 Running status runs" ); } // ----------------------------------------------------------------------- // Test 2: Concurrent limit enforcement uses batch counts // ----------------------------------------------------------------------- #[tokio::test] async fn concurrent_limit_enforcement_with_batch_counts() { let (db, _tmp) = create_test_db().await; let r1 = Uuid::new_v4(); let r2 = Uuid::new_v4(); // Create 2 routines with max_concurrent=1 (r1) and max_concurrent=2 (r2) for (routine_id, max_concurrent) in [(r1, 1), (r2, 2)] { let routine = Routine { id: routine_id, name: format!("routine-{}", routine_id), description: "Test".to_string(), user_id: "default".to_string(), enabled: true, trigger: Trigger::Cron { schedule: "* * * * *".to_string(), timezone: None, }, action: RoutineAction::Lightweight { prompt: "test".to_string(), context_paths: vec![], max_tokens: 1000, use_tools: false, max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(0), max_concurrent, dedup_window: None, }, notify: Default::default(), last_run_at: None, next_fire_at: None, run_count: 0, consecutive_failures: 0, state: serde_json::json!({}), created_at: Utc::now(), updated_at: Utc::now(), }; db.create_routine(&routine).await.expect("create routine"); } // r1: create 1 running run (will hit max_concurrent=1) let run = RoutineRun { id: Uuid::new_v4(), routine_id: r1, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: None, status: RunStatus::Running, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); // r2: create 2 running runs (will hit max_concurrent=2) for _ in 0..2 { let run = RoutineRun { id: Uuid::new_v4(), routine_id: r2, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: None, status: RunStatus::Running, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); } // Batch query should return correct counts let counts = db .count_running_routine_runs_batch(&[r1, r2]) .await .expect("batch query should work"); // Verify counts match the limits assert_eq!( counts[&r1], 1, "r1 should have 1 running (at max_concurrent=1)" ); assert_eq!( counts[&r2], 2, "r2 should have 2 running (at max_concurrent=2)" ); // Now verify the limit enforcement logic let r1_routine = db .get_routine(r1) .await .expect("get routine") .expect("routine exists"); let r2_routine = db .get_routine(r2) .await .expect("get routine") .expect("routine exists"); let r1_at_limit = counts[&r1] >= r1_routine.guardrails.max_concurrent as i64; let r2_at_limit = counts[&r2] >= r2_routine.guardrails.max_concurrent as i64; assert!(r1_at_limit, "r1 should be detected as at limit"); assert!(r2_at_limit, "r2 should be detected as at limit"); // If we add one more run to r2, it should exceed limit let run = RoutineRun { id: Uuid::new_v4(), routine_id: r2, trigger_type: "cron".to_string(), trigger_detail: None, started_at: Utc::now(), completed_at: None, status: RunStatus::Running, result_summary: None, tokens_used: None, job_id: None, created_at: Utc::now(), }; db.create_routine_run(&run).await.expect("create run"); // Re-query to get updated counts let counts = db .count_running_routine_runs_batch(&[r1, r2]) .await .expect("batch query should work"); let r2_exceeded_limit = counts[&r2] > r2_routine.guardrails.max_concurrent as i64; assert!(r2_exceeded_limit, "r2 should have exceeded its limit"); } }