rustytorch/crates/data/rtx-etl/tests/dag_test.rs

//! Comprehensive tests for DAG (Directed Acyclic Graph) functionality in RTX-ETL

use rtx_etl::{
    EtlError,
    dag::{Task, TaskGraph, TaskType},
};
use std::collections::HashMap;

#[cfg(test)]
mod dag_tests {
    use super::*;

    fn create_sample_task(id: &str, _name: &str) -> Task {
        Task::new(id.to_string(), TaskType::Transform)
    }

    #[test]
    fn test_dag_creation() {
        let dag = TaskGraph::new();
        assert_eq!(dag.tasks.len(), 0);
    }

    #[test]
    fn test_add_single_task() {
        let mut dag = TaskGraph::new();
        let task = create_sample_task("task1", "task1");

        dag.add_task(task);
        assert_eq!(dag.tasks.len(), 1);
    }

    #[test]
    fn test_add_multiple_tasks() {
        let mut dag = TaskGraph::new();

        for i in 1..=5 {
            let task = create_sample_task(&format!("task{}", i), &format!("task{}", i));
            dag.add_task(task);
        }

        assert_eq!(dag.tasks.len(), 5);
    }

    #[test]
    fn test_add_dependency() {
        let mut dag = TaskGraph::new();

        let task1 = create_sample_task("task1", "task1");
        let mut task2 = create_sample_task("task2", "task2");

        // Task2 depends on task1
        task2.dependencies.push("task1".to_string());

        dag.add_task(task1);
        dag.add_task(task2);

        assert!(
            dag.dependencies
                .get("task2")
                .unwrap()
                .contains(&"task1".to_string())
        );
    }

    #[test]
    fn test_multiple_dependencies() {
        let mut dag = TaskGraph::new();

        // Create tasks
        for i in 1..=4 {
            let task = create_sample_task(&format!("task{}", i), &format!("task{}", i));
            dag.add_task(task);
        }

        // Task 4 depends on tasks 1, 2, and 3
        let mut task4 = dag.tasks.remove("task4").unwrap();
        task4.dependencies = vec![
            "task1".to_string(),
            "task2".to_string(),
            "task3".to_string(),
        ];
        dag.add_task(task4);

        let deps = dag.dependencies.get("task4").unwrap();
        assert_eq!(deps.len(), 3);
        assert!(deps.contains(&"task1".to_string()));
        assert!(deps.contains(&"task2".to_string()));
        assert!(deps.contains(&"task3".to_string()));
    }

    #[test]
    fn test_validation_with_missing_dependency() {
        let mut dag = TaskGraph::new();

        let task1 = create_sample_task("task1", "task1");
        let mut task2 = create_sample_task("task2", "task2");

        // Task2 depends on non-existent task
        task2.dependencies.push("task999".to_string());

        dag.add_task(task1);
        dag.add_task(task2);

        let result = dag.validate();
        assert!(result.is_err());
    }

    #[test]
    fn test_remove_task() {
        let mut dag = TaskGraph::new();

        let task1 = create_sample_task("task1", "task1");
        dag.add_task(task1);

        assert_eq!(dag.tasks.len(), 1);

        let removed = dag.remove_task("task1");
        assert!(removed.is_some());
        assert_eq!(dag.tasks.len(), 0);
    }

    #[test]
    fn test_complex_dag() {
        let mut dag = TaskGraph::new();

        // Create tasks with dependencies
        let task1 = create_sample_task("task1", "task1");
        let task2 = create_sample_task("task2", "task2");
        let mut task3 = create_sample_task("task3", "task3");
        task3.dependencies = vec!["task1".to_string(), "task2".to_string()];

        let mut task4 = create_sample_task("task4", "task4");
        task4.dependencies = vec!["task2".to_string()];

        let mut task5 = create_sample_task("task5", "task5");
        task5.dependencies = vec!["task3".to_string(), "task4".to_string()];

        dag.add_task(task1);
        dag.add_task(task2);
        dag.add_task(task3);
        dag.add_task(task4);
        dag.add_task(task5);

        // Validate the graph
        let result = dag.validate();
        assert!(result.is_ok());
        assert_eq!(dag.tasks.len(), 5);
    }

    #[test]
    fn test_task_types() {
        let extract_task = Task::new("extract1".to_string(), TaskType::Extract);
        let transform_task = Task::new("transform1".to_string(), TaskType::Transform);
        let load_task = Task::new("load1".to_string(), TaskType::Load);
        let quality_task = Task::new("quality1".to_string(), TaskType::QualityCheck);

        assert_eq!(extract_task.task_type.as_str(), "extract");
        assert_eq!(transform_task.task_type.as_str(), "transform");
        assert_eq!(load_task.task_type.as_str(), "load");
        assert_eq!(quality_task.task_type.as_str(), "quality_check");
    }

    #[test]
    fn test_task_metadata() {
        let mut task = Task::new("task1".to_string(), TaskType::Transform);
        task.metadata
            .insert("key1".to_string(), "value1".to_string());
        task.metadata
            .insert("key2".to_string(), "value2".to_string());

        assert_eq!(task.metadata.get("key1").unwrap(), "value1");
        assert_eq!(task.metadata.get("key2").unwrap(), "value2");
    }

    #[test]
    fn test_task_config() {
        let mut task = Task::new("task1".to_string(), TaskType::Transform);
        task.config
            .insert("batch_size".to_string(), serde_json::json!(100));
        task.config
            .insert("timeout".to_string(), serde_json::json!(5000));

        assert_eq!(
            task.config.get("batch_size").unwrap(),
            &serde_json::json!(100)
        );
        assert_eq!(
            task.config.get("timeout").unwrap(),
            &serde_json::json!(5000)
        );
    }

    #[test]
    fn test_graph_metadata() {
        let dag = TaskGraph::new();

        assert!(!dag.metadata.graph_id.is_empty());
        assert_eq!(dag.metadata.name, "unnamed_graph");
        assert_eq!(dag.metadata.version, "1.0.0");
    }

    #[test]
    fn test_task_priority() {
        let mut high_priority = Task::new("task1".to_string(), TaskType::Extract);
        high_priority.priority = 100;

        let mut low_priority = Task::new("task2".to_string(), TaskType::Load);
        low_priority.priority = 1;

        assert!(high_priority.priority > low_priority.priority);
    }
}