Add tests for tasks

Author: CCXLV

crates/fluxqueue-worker/src/task.rs

@@ -15,11 +15,6 @@ use tokio::sync::{mpsc, oneshot};
 
 use crate::logger::Logger;
 
-type TasksAndContexts = (
-    HashMap<String, Arc<Py<PyAny>>>,
-    HashMap<String, Arc<Py<PyAny>>>,
-);
-
 #[derive(Debug)]
 pub struct TaskRegistry {
     tasks: Arc<RwLock<HashMap<String, Arc<Py<PyAny>>>>>,
@@ -28,75 +23,7 @@ pub struct TaskRegistry {
 
 impl TaskRegistry {
     pub fn new(module_path: &str, queue_name: &str) -> Result<Self> {
-        let script = include_str!("../scripts/get_registry.py");
-        let script_cstr = CString::new(script)?;
-        let filename = CString::new("get_registry.py")?;
-        let module_name = CString::new("get_registry")?;
-
-        let full_current_dir = std::env::current_dir().unwrap();
-        let full_module_path = full_current_dir.join(module_path);
-        let clean_module_path = normalize_path(&full_module_path);
-        let project_root = full_current_dir
-            .ancestors()
-            .find(|p| p.join("tests").exists())
-            .unwrap_or(&full_current_dir);
-        let real_module_path = path_to_module_path(project_root, &clean_module_path);
-
-        if !clean_module_path.exists() || real_module_path.is_none() {
-            return Err(anyhow!(
-                "Tasks module path {:?} doesn't exist.",
-                clean_module_path
-            ));
-        }
-
-        let real_module_path = real_module_path.unwrap();
-        let module_dir = project_root.to_string_lossy().to_string();
-
-        let (tasks, contexts) = Python::attach(|py| -> Result<TasksAndContexts> {
-            let module = PyModule::from_code(
-                py,
-                script_cstr.as_c_str(),
-                filename.as_c_str(),
-                module_name.as_c_str(),
-            )
-            .map_err(|e| anyhow!("Failed to import python module: {}", e))?;
-
-            let registry: Bound<'_, PyDict> = module
-                .getattr("get_registry")
-                .map_err(|e| anyhow!("Failed to get 'get_registry' script: {}", e))?
-                .call1((real_module_path, queue_name, module_dir))
-                .map_err(|e| anyhow!("Failed to get tasks: {}", e))?
-                .cast_into::<PyDict>()
-                .map_err(|_| anyhow!("Failed to cast result to a Python Dictionary"))?;
-
-            let tasks: HashMap<String, Arc<Py<PyAny>>> = registry
-                .get_item("tasks")?
-                .expect("tasks missing")
-                .cast::<PyDict>()
-                .map_err(|e| anyhow!("tasks is not a dict: {}", e))?
-                .iter()
-                .filter_map(|(key, value)| {
-                    let name: String = key.extract().ok()?;
-                    let func: Py<PyAny> = value.unbind();
-                    Some((name, Arc::new(func)))
-                })
-                .collect();
-
-            let contexts: HashMap<String, Arc<Py<PyAny>>> = registry
-                .get_item("contexts")?
-                .expect("contexts missing")
-                .cast::<PyDict>()
-                .map_err(|e| anyhow!("contexts is not a dict: {}", e))?
-                .iter()
-                .filter_map(|(key, value): (Bound<PyAny>, Bound<PyAny>)| {
-                    let name: String = key.extract().ok()?;
-                    let func: Py<PyAny> = value.unbind();
-                    Some((name, Arc::new(func)))
-                })
-                .collect();
-
-            Ok((tasks, contexts))
-        })?;
+        let (tasks, contexts) = get_registry(module_path, queue_name)?;
 
         Ok(Self {
             tasks: Arc::new(RwLock::new(tasks)),
@@ -276,6 +203,85 @@ async fn run_task(
     Ok(())
 }
 
+type TasksAndContexts = (
+    HashMap<String, Arc<Py<PyAny>>>,
+    HashMap<String, Arc<Py<PyAny>>>,
+);
+
+fn get_registry(module_path: &str, queue_name: &str) -> Result<TasksAndContexts> {
+    let script = include_str!("../scripts/get_registry.py");
+    let script_cstr = CString::new(script)?;
+    let filename = CString::new("get_registry.py")?;
+    let module_name = CString::new("get_registry")?;
+
+    let full_current_dir = std::env::current_dir().unwrap();
+    let full_module_path = full_current_dir.join(module_path);
+    let clean_module_path = normalize_path(&full_module_path);
+    let project_root = full_current_dir
+        .ancestors()
+        .find(|p| p.join("tests").exists())
+        .unwrap_or(&full_current_dir);
+    let real_module_path = path_to_module_path(project_root, &clean_module_path);
+
+    if !clean_module_path.exists() || real_module_path.is_none() {
+        return Err(anyhow!(
+            "Tasks module path {:?} doesn't exist.",
+            clean_module_path
+        ));
+    }
+
+    let real_module_path = real_module_path.unwrap();
+    let module_dir = project_root.to_string_lossy().to_string();
+
+    let result = Python::attach(|py| -> Result<TasksAndContexts> {
+        let module = PyModule::from_code(
+            py,
+            script_cstr.as_c_str(),
+            filename.as_c_str(),
+            module_name.as_c_str(),
+        )
+        .map_err(|e| anyhow!("Failed to import python module: {}", e))?;
+
+        let registry: Bound<'_, PyDict> = module
+            .getattr("get_registry")
+            .map_err(|e| anyhow!("Failed to get 'get_registry' script: {}", e))?
+            .call1((real_module_path, queue_name, module_dir))
+            .map_err(|e| anyhow!("Failed to get tasks: {}", e))?
+            .cast_into::<PyDict>()
+            .map_err(|_| anyhow!("Failed to cast result to a Python Dictionary"))?;
+
+        let tasks: HashMap<String, Arc<Py<PyAny>>> = registry
+            .get_item("tasks")?
+            .expect("tasks missing")
+            .cast::<PyDict>()
+            .map_err(|e| anyhow!("tasks is not a dict: {}", e))?
+            .iter()
+            .filter_map(|(key, value)| {
+                let name: String = key.extract().ok()?;
+                let func: Py<PyAny> = value.unbind();
+                Some((name, Arc::new(func)))
+            })
+            .collect();
+
+        let contexts: HashMap<String, Arc<Py<PyAny>>> = registry
+            .get_item("contexts")?
+            .expect("contexts missing")
+            .cast::<PyDict>()
+            .map_err(|e| anyhow!("contexts is not a dict: {}", e))?
+            .iter()
+            .filter_map(|(key, value): (Bound<PyAny>, Bound<PyAny>)| {
+                let name: String = key.extract().ok()?;
+                let func: Py<PyAny> = value.unbind();
+                Some((name, Arc::new(func)))
+            })
+            .collect();
+
+        Ok((tasks, contexts))
+    })?;
+
+    Ok(result)
+}
+
 fn normalize_path(path: &Path) -> PathBuf {
     let mut components = Vec::new();
     for comp in path.components() {
@@ -306,3 +312,84 @@ fn path_to_module_path(current_dir: &Path, target_path: &Path) -> Option<String>
 
     Some(components.join("."))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_path_to_module_path() -> Result<()> {
+        let current_dir = Path::new("project");
+        let tasks_path = Path::new("../project/tasks.py");
+        let normalized_path = normalize_path(tasks_path);
+        let module_path = path_to_module_path(current_dir, &normalized_path);
+        let expected_path = Path::new("project/tasks.py");
+
+        assert_eq!(normalized_path, expected_path);
+        assert_eq!(module_path, Some("tasks".to_string()));
+
+        Ok(())
+    }
+
+    fn get_test_module_path(filename: &str) -> String {
+        let current_dir = std::env::current_dir().unwrap();
+        let test_module_path = current_dir.join("tests").join(filename);
+        test_module_path.to_str().unwrap().to_string()
+    }
+
+    #[test]
+    fn test_get_task_functions_valid_module() -> Result<()> {
+        let module_path_str = get_test_module_path("test_tasks_module.py");
+        let (tasks, _) = get_registry(&module_path_str, "default")?;
+
+        assert_eq!(tasks.len(), 3);
+
+        let task_names: Vec<String> = tasks.iter().map(|(name, _)| name.clone()).collect();
+        assert!(task_names.contains(&"task-1".to_string()));
+        assert!(task_names.contains(&"task-2".to_string()));
+        assert!(task_names.contains(&"async-task".to_string()));
+
+        assert!(!task_names.contains(&"high-priority-task".to_string()));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_task_functions_different_queue() -> Result<()> {
+        let module_path_str = get_test_module_path("test_tasks_module.py");
+        let (tasks, _) = get_registry(&module_path_str, "high-priority")?;
+
+        let task_names: Vec<String> = tasks.iter().map(|(name, _)| name.clone()).collect();
+        assert_eq!(tasks.len(), 1);
+        assert!(task_names.contains(&"high-priority-task".to_string()));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_task_functions_empty_module() -> Result<()> {
+        let module_path_str = get_test_module_path("test_tasks_empty.py");
+        let (tasks, _) = get_registry(&module_path_str, "default")?;
+
+        assert_eq!(tasks.len(), 0);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_task_functions_duplicate_names() {
+        let module_path_str = get_test_module_path("test_tasks_duplicate.py");
+
+        let result = get_registry(&module_path_str, "default");
+        assert!(result.is_err());
+
+        let error_msg = result.unwrap_err().to_string();
+        assert!(error_msg.contains("duplicated") || error_msg.contains("duplicate"));
+    }
+
+    #[test]
+    fn test_get_task_functions_invalid_path() {
+        let result = get_registry("nonexistent/path/to/module.py", "default");
+        assert!(result.is_err());
+    }
+}