replace default variables node_address, task_id on orchestrator (#457)

Author: JannikSt

crates/orchestrator/src/scheduler/mod.rs

@@ -10,7 +10,6 @@ pub struct Scheduler {
     store_context: Arc<StoreContext>,
     plugins: Vec<Box<dyn SchedulerPlugin>>,
 }
-
 impl Scheduler {
     pub fn new(store_context: Arc<StoreContext>, plugins: Vec<Box<dyn SchedulerPlugin>>) -> Self {
         let mut plugins = plugins;
@@ -33,7 +32,28 @@ impl Scheduler {
         }
 
         if !all_tasks.is_empty() {
-            return Ok(Some(all_tasks[0].clone()));
+            let mut task = all_tasks[0].clone();
+
+            // Replace variables in env_vars
+            if let Some(env_vars) = &mut task.env_vars {
+                for (_, value) in env_vars.iter_mut() {
+                    let new_value = value
+                        .replace("${TASK_ID}", &task.id.to_string())
+                        .replace("${NODE_ADDRESS}", &node_address.to_string());
+                    *value = new_value;
+                }
+            }
+
+            // Replace variables in args
+            if let Some(args) = &mut task.args {
+                for arg in args.iter_mut() {
+                    *arg = arg
+                        .replace("${TASK_ID}", &task.id.to_string())
+                        .replace("${NODE_ADDRESS}", &node_address.to_string());
+                }
+            }
+
+            return Ok(Some(task));
         }
 
         Ok(None)
@@ -43,6 +63,7 @@ impl Scheduler {
 #[cfg(test)]
 mod tests {
     use shared::models::task::TaskState;
+    use std::collections::HashMap;
     use uuid::Uuid;
 
     use crate::api::tests::helper::create_test_app_state;
@@ -68,4 +89,52 @@ mod tests {
         let task_for_node = scheduler.get_task_for_node(Address::ZERO).await.unwrap();
         assert_eq!(task_for_node, Some(task));
     }
+
+    #[tokio::test]
+    async fn test_variable_replacement() {
+        let state = create_test_app_state().await;
+        let scheduler = Scheduler::new(state.store_context.clone(), vec![]);
+        let node_address = Address::from([1u8; 20]);
+
+        let mut env_vars = HashMap::new();
+        env_vars.insert("TASK_ID_VAR".to_string(), "task-${TASK_ID}".to_string());
+        env_vars.insert("NODE_VAR".to_string(), "node-${NODE_ADDRESS}".to_string());
+
+        let task = Task {
+            id: Uuid::new_v4(),
+            image: "image".to_string(),
+            name: "name".to_string(),
+            state: TaskState::PENDING,
+            created_at: 1,
+            env_vars: Some(env_vars),
+            args: Some(vec![
+                "--task=${TASK_ID}".to_string(),
+                "--node=${NODE_ADDRESS}".to_string(),
+            ]),
+            ..Default::default()
+        };
+
+        let _ = state.store_context.task_store.add_task(task.clone()).await;
+
+        let result = scheduler.get_task_for_node(node_address).await.unwrap();
+        assert!(result.is_some());
+
+        let returned_task = result.unwrap();
+
+        // Check env vars replacement
+        let env_vars = returned_task.env_vars.unwrap();
+        assert_eq!(
+            env_vars.get("TASK_ID_VAR").unwrap(),
+            &format!("task-{}", task.id)
+        );
+        assert_eq!(
+            env_vars.get("NODE_VAR").unwrap(),
+            &format!("node-{}", node_address)
+        );
+
+        // Check args replacement
+        let args = returned_task.args.unwrap();
+        assert_eq!(args[0], format!("--task={}", task.id));
+        assert_eq!(args[1], format!("--node={}", node_address));
+    }
 }