Skip to main content

ironflow_api/entities/
worker_capabilities.rs

1//! Worker capabilities sent as query parameters on the internal pick and
2//! pending-count routes.
3
4use ironflow_store::entities::{WorkerCapabilities, normalize_worker_tags, validate_worker_tags};
5
6use crate::error::ApiError;
7
8/// Parse the `workflows` and `tags` query parameters a worker sends.
9///
10/// Both are comma-separated lists. When neither is present the worker
11/// predates routing: `None` is returned and it takes every run. When only
12/// `tags` is present the worker accepts any workflow. An empty `tags` value
13/// means a worker carrying no tag: it only takes runs that require none.
14///
15/// # Errors
16///
17/// Returns [`ApiError::BadRequest`] when a tag is invalid (see
18/// [`validate_worker_tags`]).
19///
20/// # Examples
21///
22/// ```
23/// use ironflow_api::entities::parse_worker_capabilities;
24///
25/// # fn example() -> Result<(), ironflow_api::error::ApiError> {
26/// assert!(parse_worker_capabilities(None, None)?.is_none());
27///
28/// let caps = parse_worker_capabilities(Some("deploy,build"), Some("region:eu,gpu"))?
29///     .expect("capabilities were sent");
30/// assert_eq!(caps.workflows, Some(vec!["build".to_string(), "deploy".to_string()]));
31/// assert_eq!(caps.tags, vec!["gpu".to_string(), "region:eu".to_string()]);
32/// # Ok(())
33/// # }
34/// ```
35pub fn parse_worker_capabilities(
36    workflows: Option<&str>,
37    tags: Option<&str>,
38) -> Result<Option<WorkerCapabilities>, ApiError> {
39    if workflows.is_none() && tags.is_none() {
40        return Ok(None);
41    }
42
43    let workflows = workflows.map(|raw| normalize_worker_tags(split_list(raw)));
44    let tags: Vec<String> = tags.map(split_list).unwrap_or_default();
45    validate_worker_tags(&tags).map_err(|e| ApiError::BadRequest(e.to_string()))?;
46
47    Ok(Some(WorkerCapabilities::new(
48        workflows,
49        normalize_worker_tags(tags),
50    )))
51}
52
53/// Split a comma-separated list, dropping blank entries.
54fn split_list(raw: &str) -> Vec<String> {
55    raw.split(',')
56        .map(str::trim)
57        .filter(|entry| !entry.is_empty())
58        .map(str::to_string)
59        .collect()
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    fn strings(values: &[&str]) -> Vec<String> {
67        values.iter().map(|v| (*v).to_string()).collect()
68    }
69
70    #[test]
71    fn absent_params_mean_a_legacy_worker() {
72        assert!(parse_worker_capabilities(None, None).unwrap().is_none());
73    }
74
75    #[test]
76    fn both_params_are_parsed() {
77        let caps = parse_worker_capabilities(Some("deploy, build"), Some("gpu,region:eu"))
78            .unwrap()
79            .unwrap();
80        assert_eq!(caps.workflows, Some(strings(&["build", "deploy"])));
81        assert_eq!(caps.tags, strings(&["gpu", "region:eu"]));
82    }
83
84    #[test]
85    fn tags_without_workflows_accept_any_workflow() {
86        let caps = parse_worker_capabilities(None, Some("gpu"))
87            .unwrap()
88            .unwrap();
89        assert!(caps.workflows.is_none());
90        assert_eq!(caps.tags, strings(&["gpu"]));
91    }
92
93    #[test]
94    fn empty_tags_mean_a_worker_without_tags() {
95        let caps = parse_worker_capabilities(Some("deploy"), Some(""))
96            .unwrap()
97            .unwrap();
98        assert_eq!(caps.workflows, Some(strings(&["deploy"])));
99        assert!(caps.tags.is_empty());
100    }
101
102    #[test]
103    fn empty_workflows_take_no_workflow() {
104        let caps = parse_worker_capabilities(Some(""), None).unwrap().unwrap();
105        assert_eq!(caps.workflows, Some(Vec::new()));
106        assert!(caps.tags.is_empty());
107    }
108
109    #[test]
110    fn duplicate_tags_are_dropped() {
111        let caps = parse_worker_capabilities(None, Some("gpu,gpu, gpu,,arm"))
112            .unwrap()
113            .unwrap();
114        assert_eq!(caps.tags, strings(&["arm", "gpu"]));
115    }
116
117    #[test]
118    fn invalid_tag_is_a_bad_request() {
119        let err = parse_worker_capabilities(None, Some("gpu,two words")).unwrap_err();
120        assert!(matches!(err, ApiError::BadRequest(ref msg) if msg.contains("two words")));
121    }
122
123    #[test]
124    fn too_long_tag_is_a_bad_request() {
125        let long = "a".repeat(65);
126        let err = parse_worker_capabilities(None, Some(&long)).unwrap_err();
127        assert!(matches!(err, ApiError::BadRequest(_)));
128    }
129}