1use chrono::Utc;
7use std::collections::HashMap;
8
9use tatara_core::cluster::types::{NodeMeta, NodeRoles};
10use tatara_core::domain::allocation::Allocation;
11use tatara_core::domain::job::*;
12use tatara_core::domain::release::{Release, ReleaseStatus};
13
14pub fn job(id: &str) -> Job {
16 job_with_group(id, "main", 1, 500, 256)
17}
18
19pub fn job_with_group(id: &str, group_name: &str, count: u32, cpu_mhz: u64, memory_mb: u64) -> Job {
21 Job {
22 id: id.to_string(),
23 version: 1,
24 job_type: JobType::Service,
25 status: JobStatus::Pending,
26 submitted_at: Utc::now(),
27 groups: vec![TaskGroup {
28 name: group_name.to_string(),
29 count,
30 tasks: vec![task("app", cpu_mhz, memory_mb)],
31 restart_policy: RestartPolicy::default(),
32 resources: Resources { cpu_mhz, memory_mb },
33 network: None,
34 secrets: vec![],
35 volumes: vec![],
36 service_name: None,
37 }],
38 constraints: vec![],
39 meta: HashMap::new(),
40 spec_hash: None,
41 }
42}
43
44pub fn job_spec(id: &str) -> JobSpec {
46 job_spec_with_group(id, "main", 1, 500, 256)
47}
48
49pub fn job_spec_with_group(
51 id: &str,
52 group_name: &str,
53 count: u32,
54 cpu_mhz: u64,
55 memory_mb: u64,
56) -> JobSpec {
57 JobSpec {
58 id: id.to_string(),
59 job_type: JobType::Service,
60 groups: vec![TaskGroup {
61 name: group_name.to_string(),
62 count,
63 tasks: vec![task("app", cpu_mhz, memory_mb)],
64 restart_policy: RestartPolicy::default(),
65 resources: Resources { cpu_mhz, memory_mb },
66 network: None,
67 secrets: vec![],
68 volumes: vec![],
69 service_name: None,
70 }],
71 constraints: vec![],
72 meta: HashMap::new(),
73 }
74}
75
76pub fn batch_job_spec(id: &str, cpu_mhz: u64, memory_mb: u64) -> JobSpec {
78 JobSpec {
79 id: id.to_string(),
80 job_type: JobType::Batch,
81 groups: vec![TaskGroup {
82 name: "main".to_string(),
83 count: 1,
84 tasks: vec![task("worker", cpu_mhz, memory_mb)],
85 restart_policy: RestartPolicy {
86 mode: RestartMode::Never,
87 ..Default::default()
88 },
89 resources: Resources { cpu_mhz, memory_mb },
90 network: None,
91 secrets: vec![],
92 volumes: vec![],
93 service_name: None,
94 }],
95 constraints: vec![],
96 meta: HashMap::new(),
97 }
98}
99
100pub fn constrained_job_spec(id: &str, constraints: Vec<Constraint>) -> JobSpec {
102 let mut spec = job_spec(id);
103 spec.constraints = constraints;
104 spec
105}
106
107pub fn forge_job_spec(name: &str, flake_ref: &str) -> JobSpec {
109 JobSpec {
110 id: name.to_string(),
111 job_type: JobType::Service,
112 groups: vec![TaskGroup {
113 name: "main".to_string(),
114 count: 1,
115 tasks: vec![Task {
116 name: "app".to_string(),
117 driver: DriverType::Nix,
118 config: TaskConfig::Nix {
119 flake_ref: flake_ref.to_string(),
120 args: vec![],
121 },
122 env: HashMap::new(),
123 resources: Resources {
124 cpu_mhz: 500,
125 memory_mb: 256,
126 },
127 health_checks: vec![],
128 volume_claims: vec![],
129 }],
130 restart_policy: RestartPolicy::default(),
131 resources: Resources {
132 cpu_mhz: 500,
133 memory_mb: 256,
134 },
135 network: None,
136 secrets: vec![],
137 volumes: vec![],
138 service_name: None,
139 }],
140 constraints: vec![],
141 meta: {
142 let mut m = HashMap::new();
143 m.insert("forge".to_string(), "true".to_string());
144 m.insert("flake_ref".to_string(), flake_ref.to_string());
145 m
146 },
147 }
148}
149
150pub fn task(name: &str, cpu_mhz: u64, memory_mb: u64) -> Task {
152 Task {
153 name: name.to_string(),
154 driver: DriverType::Exec,
155 config: TaskConfig::Exec {
156 command: "echo".to_string(),
157 args: vec!["hello".to_string()],
158 working_dir: None,
159 },
160 env: HashMap::new(),
161 resources: Resources { cpu_mhz, memory_mb },
162 health_checks: vec![],
163 volume_claims: vec![],
164 }
165}
166
167pub fn node_meta(node_id: u64, hostname: &str, cpu_mhz: u64, memory_mb: u64) -> NodeMeta {
169 NodeMeta {
170 node_id,
171 hostname: hostname.to_string(),
172 http_addr: format!("127.0.0.1:{}", 4646 + node_id),
173 gossip_addr: format!("127.0.0.1:{}", 4648 + node_id),
174 raft_addr: format!("127.0.0.1:{}", 4649 + node_id),
175 os: "linux".to_string(),
176 arch: "x86_64".to_string(),
177 roles: NodeRoles::default(),
178 drivers: vec![DriverType::Exec, DriverType::Nix],
179 total_resources: Resources { cpu_mhz, memory_mb },
180 available_resources: Resources { cpu_mhz, memory_mb },
181 allocations_running: 0,
182 joined_at: Utc::now(),
183 version: "0.2.0".to_string(),
184 eligible: true,
185 wireguard_pubkey: None,
186 tunnel_address: None,
187 }
188}
189
190pub fn node_meta_with_os(
192 node_id: u64,
193 hostname: &str,
194 os: &str,
195 arch: &str,
196 cpu_mhz: u64,
197 memory_mb: u64,
198) -> NodeMeta {
199 let mut meta = node_meta(node_id, hostname, cpu_mhz, memory_mb);
200 meta.os = os.to_string();
201 meta.arch = arch.to_string();
202 meta
203}
204
205pub fn allocation(job_id: &str, group_name: &str, node_id: &str) -> Allocation {
207 Allocation::new(
208 job_id.to_string(),
209 group_name.to_string(),
210 node_id.to_string(),
211 vec!["app".to_string()],
212 )
213}
214
215pub fn constraint(attribute: &str, operator: &str, value: &str) -> Constraint {
217 Constraint {
218 attribute: attribute.to_string(),
219 operator: operator.to_string(),
220 value: value.to_string(),
221 }
222}
223
224pub fn release(name: &str, flake_ref: &str, job_id: &str) -> Release {
226 let mut r = Release::new(name.to_string(), flake_ref.to_string(), job_id.to_string());
227 r.status = ReleaseStatus::Active;
228 r
229}
230
231pub fn pending_release(name: &str, flake_ref: &str, job_id: &str) -> Release {
233 Release::new(name.to_string(), flake_ref.to_string(), job_id.to_string())
234}