1use std::collections::HashMap;
2use std::marker::PhantomData;
3
4use crate::api::{IngressConfig, JobSubmitRequestItem, ServiceConfig};
5use crate::k8s::io::api::core::v1::PodSpec;
6
7#[doc(hidden)]
9pub struct NoPodSpec;
10#[doc(hidden)]
12pub struct HasPodSpec;
13
14pub struct JobRequestItemBuilder<S> {
32 priority: f64,
33 namespace: String,
34 client_id: String,
35 labels: HashMap<String, String>,
36 annotations: HashMap<String, String>,
37 scheduler: String,
38 ingress: Vec<IngressConfig>,
39 services: Vec<ServiceConfig>,
40 pod_specs: Vec<PodSpec>,
41 _state: PhantomData<S>,
42}
43
44impl JobRequestItemBuilder<NoPodSpec> {
45 pub fn new() -> Self {
52 Self {
53 priority: 0.0,
54 namespace: String::new(),
55 client_id: String::new(),
56 labels: HashMap::new(),
57 annotations: HashMap::new(),
58 scheduler: String::new(),
59 ingress: Vec::new(),
60 services: Vec::new(),
61 pod_specs: Vec::new(),
62 _state: PhantomData,
63 }
64 }
65}
66
67impl Default for JobRequestItemBuilder<NoPodSpec> {
68 fn default() -> Self {
69 Self::new()
70 }
71}
72
73impl<S> JobRequestItemBuilder<S> {
75 #[must_use]
79 pub fn priority(mut self, p: f64) -> Self {
80 self.priority = p;
81 self
82 }
83
84 #[must_use]
89 pub fn namespace(mut self, ns: impl Into<String>) -> Self {
90 self.namespace = ns.into();
91 self
92 }
93
94 #[must_use]
100 pub fn client_id(mut self, id: impl Into<String>) -> Self {
101 self.client_id = id.into();
102 self
103 }
104
105 #[must_use]
107 pub fn labels(mut self, l: HashMap<String, String>) -> Self {
108 self.labels = l;
109 self
110 }
111
112 #[must_use]
114 pub fn label(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
115 self.labels.insert(k.into(), v.into());
116 self
117 }
118
119 #[must_use]
121 pub fn annotations(mut self, a: HashMap<String, String>) -> Self {
122 self.annotations = a;
123 self
124 }
125
126 #[must_use]
128 pub fn annotation(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
129 self.annotations.insert(k.into(), v.into());
130 self
131 }
132
133 #[must_use]
135 pub fn scheduler(mut self, s: impl Into<String>) -> Self {
136 self.scheduler = s.into();
137 self
138 }
139
140 #[must_use]
142 pub fn ingress(mut self, i: Vec<IngressConfig>) -> Self {
143 self.ingress = i;
144 self
145 }
146
147 #[must_use]
149 pub fn add_ingress(mut self, i: IngressConfig) -> Self {
150 self.ingress.push(i);
151 self
152 }
153
154 #[must_use]
156 pub fn services(mut self, s: Vec<ServiceConfig>) -> Self {
157 self.services = s;
158 self
159 }
160
161 #[must_use]
163 pub fn add_service(mut self, s: ServiceConfig) -> Self {
164 self.services.push(s);
165 self
166 }
167
168 #[must_use]
174 pub fn pod_spec(self, spec: PodSpec) -> JobRequestItemBuilder<HasPodSpec> {
175 self.pod_specs(vec![spec])
176 }
177
178 #[must_use]
184 pub fn pod_specs(self, specs: Vec<PodSpec>) -> JobRequestItemBuilder<HasPodSpec> {
185 JobRequestItemBuilder {
186 priority: self.priority,
187 namespace: self.namespace,
188 client_id: self.client_id,
189 labels: self.labels,
190 annotations: self.annotations,
191 scheduler: self.scheduler,
192 ingress: self.ingress,
193 services: self.services,
194 pod_specs: specs,
195 _state: PhantomData,
196 }
197 }
198}
199
200impl JobRequestItemBuilder<HasPodSpec> {
202 #[must_use]
210 pub fn build(self) -> JobSubmitRequestItem {
211 #[allow(deprecated)]
212 JobSubmitRequestItem {
213 priority: self.priority,
214 namespace: self.namespace,
215 client_id: self.client_id,
216 labels: self.labels,
217 annotations: self.annotations,
218 pod_specs: self.pod_specs,
219 ingress: self.ingress,
220 services: self.services,
221 scheduler: self.scheduler,
222 pod_spec: None,
224 required_node_labels: HashMap::new(),
225 }
226 }
227}
228
229#[cfg(test)]
230mod tests {
231 use super::*;
232
233 fn minimal_pod_spec() -> PodSpec {
234 PodSpec {
235 containers: vec![],
236 ..Default::default()
237 }
238 }
239
240 #[test]
241 fn builder_with_pod_specs_builds_correctly() {
242 let spec = minimal_pod_spec();
243 let item = JobRequestItemBuilder::new()
244 .namespace("default")
245 .priority(1.0)
246 .pod_specs(vec![spec])
247 .build();
248
249 assert_eq!(item.namespace, "default");
250 assert_eq!(item.priority, 1.0);
251 assert_eq!(item.pod_specs.len(), 1);
252 }
253
254 #[test]
255 fn pod_spec_singular_shorthand() {
256 let item = JobRequestItemBuilder::new()
257 .pod_spec(minimal_pod_spec())
258 .build();
259 assert_eq!(item.pod_specs.len(), 1);
260 }
261
262 #[test]
263 fn pod_spec_called_twice_replaces_previous() {
264 let spec_a = minimal_pod_spec();
265 let mut spec_b = minimal_pod_spec();
266 spec_b.restart_policy = Some("Never".to_string());
267
268 let item = JobRequestItemBuilder::new()
269 .pod_spec(spec_a)
270 .pod_spec(spec_b)
271 .build();
272
273 assert_eq!(item.pod_specs.len(), 1);
274 assert_eq!(item.pod_specs[0].restart_policy.as_deref(), Some("Never"));
275 }
276
277 #[test]
278 fn label_and_annotation_helpers() {
279 let item = JobRequestItemBuilder::new()
280 .label("app", "my-app")
281 .label("env", "prod")
282 .annotation("owner", "team-a")
283 .pod_spec(minimal_pod_spec())
284 .build();
285
286 assert_eq!(item.labels.get("app").map(String::as_str), Some("my-app"));
287 assert_eq!(item.labels.get("env").map(String::as_str), Some("prod"));
288 assert_eq!(
289 item.annotations.get("owner").map(String::as_str),
290 Some("team-a")
291 );
292 }
293
294 #[test]
295 fn optional_fields_default_to_empty() {
296 let item = JobRequestItemBuilder::new()
297 .pod_specs(vec![minimal_pod_spec()])
298 .build();
299
300 assert_eq!(item.priority, 0.0);
301 assert!(item.namespace.is_empty());
302 assert!(item.client_id.is_empty());
303 assert!(item.labels.is_empty());
304 assert!(item.annotations.is_empty());
305 assert!(item.scheduler.is_empty());
306 assert!(item.ingress.is_empty());
307 assert!(item.services.is_empty());
308 }
309}