capsule_core/wasm/utilities/
task_config.rs

1use serde::Deserialize;
2
3use crate::wasm::execution_policy::{Compute, ExecutionPolicy};
4
5#[derive(Debug, Deserialize, Default)]
6pub struct TaskConfig {
7    name: Option<String>,
8    compute: Option<String>,
9    ram: Option<String>,
10    timeout: Option<String>,
11
12    #[serde(alias = "maxRetries")]
13    max_retries: Option<u64>,
14}
15
16impl TaskConfig {
17    pub fn to_execution_policy(&self) -> ExecutionPolicy {
18        let compute = self
19            .compute
20            .as_ref()
21            .map(|c| match c.to_uppercase().as_str() {
22                "LOW" => Compute::Low,
23                "MEDIUM" => Compute::Medium,
24                "HIGH" => Compute::High,
25                _ => c
26                    .parse::<u64>()
27                    .map(Compute::Custom)
28                    .unwrap_or(Compute::Medium),
29            });
30
31        let ram = self.ram.as_ref().and_then(|r| Self::parse_ram_string(r));
32
33        ExecutionPolicy::new()
34            .name(self.name.clone())
35            .compute(compute)
36            .ram(ram)
37            .timeout(self.timeout.clone())
38            .max_retries(self.max_retries)
39    }
40
41    pub fn parse_ram_string(s: &str) -> Option<u64> {
42        let s = s.trim().to_uppercase();
43        if s.ends_with("GB") {
44            s.trim_end_matches("GB")
45                .trim()
46                .parse::<u64>()
47                .ok()
48                .map(|v| v * 1024 * 1024 * 1024)
49        } else if s.ends_with("MB") {
50            s.trim_end_matches("MB")
51                .trim()
52                .parse::<u64>()
53                .ok()
54                .map(|v| v * 1024 * 1024)
55        } else if s.ends_with("KB") {
56            s.trim_end_matches("KB")
57                .trim()
58                .parse::<u64>()
59                .ok()
60                .map(|v| v * 1024)
61        } else {
62            s.parse::<u64>().ok()
63        }
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn test_parse_ram_string() {
73        assert_eq!(
74            TaskConfig::parse_ram_string("2GB"),
75            Some(2 * 1024 * 1024 * 1024)
76        );
77        assert_eq!(
78            TaskConfig::parse_ram_string("1 GB"),
79            Some(1024 * 1024 * 1024)
80        );
81        assert_eq!(
82            TaskConfig::parse_ram_string("4gb"),
83            Some(4 * 1024 * 1024 * 1024)
84        );
85
86        assert_eq!(
87            TaskConfig::parse_ram_string("512MB"),
88            Some(512 * 1024 * 1024)
89        );
90        assert_eq!(
91            TaskConfig::parse_ram_string("256 MB"),
92            Some(256 * 1024 * 1024)
93        );
94        assert_eq!(
95            TaskConfig::parse_ram_string("128mb"),
96            Some(128 * 1024 * 1024)
97        );
98
99        assert_eq!(TaskConfig::parse_ram_string("1024KB"), Some(1024 * 1024));
100        assert_eq!(TaskConfig::parse_ram_string("512 KB"), Some(512 * 1024));
101        assert_eq!(TaskConfig::parse_ram_string("256kb"), Some(256 * 1024));
102
103        assert_eq!(TaskConfig::parse_ram_string("1024"), Some(1024));
104        assert_eq!(TaskConfig::parse_ram_string("512"), Some(512));
105    }
106
107    #[test]
108    fn test_parse_ram_string_invalid() {
109        assert_eq!(TaskConfig::parse_ram_string("invalid"), None);
110        assert_eq!(TaskConfig::parse_ram_string(""), None);
111        assert_eq!(TaskConfig::parse_ram_string("GB"), None);
112    }
113
114    #[test]
115    fn test_to_execution_policy_default() {
116        let config = TaskConfig::default();
117        let policy = config.to_execution_policy();
118
119        assert_eq!(policy.name, "default");
120        assert_eq!(policy.compute, Compute::Medium);
121        assert_eq!(policy.ram, None);
122        assert_eq!(policy.timeout, None);
123        assert_eq!(policy.max_retries, 0);
124    }
125
126    #[test]
127    fn test_to_execution_policy_with_values() {
128        let config = TaskConfig {
129            name: Some("test_task".to_string()),
130            compute: Some("HIGH".to_string()),
131            ram: Some("2GB".to_string()),
132            timeout: Some("30s".to_string()),
133            max_retries: Some(3),
134        };
135
136        let policy = config.to_execution_policy();
137
138        assert_eq!(policy.name, "test_task");
139        assert_eq!(policy.compute, Compute::High);
140        assert_eq!(policy.ram, Some(2 * 1024 * 1024 * 1024));
141        assert_eq!(policy.timeout, Some("30s".to_string()));
142        assert_eq!(policy.max_retries, 3);
143    }
144
145    #[test]
146    fn test_to_execution_policy_compute_variants() {
147        let low = TaskConfig {
148            compute: Some("LOW".to_string()),
149            ..Default::default()
150        };
151        assert_eq!(low.to_execution_policy().compute, Compute::Low);
152
153        let medium = TaskConfig {
154            compute: Some("MEDIUM".to_string()),
155            ..Default::default()
156        };
157        assert_eq!(medium.to_execution_policy().compute, Compute::Medium);
158
159        let high = TaskConfig {
160            compute: Some("HIGH".to_string()),
161            ..Default::default()
162        };
163        assert_eq!(high.to_execution_policy().compute, Compute::High);
164
165        let invalid = TaskConfig {
166            compute: Some("INVALID".to_string()),
167            ..Default::default()
168        };
169        assert_eq!(invalid.to_execution_policy().compute, Compute::Medium);
170    }
171}
capsule_core/wasm/utilities/task_config.rs

capsule_core/wasm/utilities/
task_config.rs
