Skip to main content

tatara_testing/
fixtures.rs

1//! Test fixture builders for creating domain objects.
2//!
3//! These functions provide convenient constructors for test data,
4//! using sensible defaults while allowing customization.
5
6use 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
14/// Create a minimal job with sensible defaults.
15pub fn job(id: &str) -> Job {
16    job_with_group(id, "main", 1, 500, 256)
17}
18
19/// Create a job with a specific task group configuration.
20pub 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
44/// Create a job spec (pre-submission form).
45pub fn job_spec(id: &str) -> JobSpec {
46    job_spec_with_group(id, "main", 1, 500, 256)
47}
48
49/// Create a job spec with specific group configuration.
50pub 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
76/// Create a batch job spec.
77pub 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
100/// Create a job spec with constraints.
101pub 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
107/// Create a job spec that looks like a forge-deployed workload.
108pub 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
150/// Create an exec task with given resource requirements.
151pub 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
167/// Create a node metadata entry.
168pub 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
190/// Create a node metadata entry with custom attributes baked into os/arch.
191pub 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
205/// Create an allocation for a job on a node.
206pub 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
215/// Create a constraint.
216pub 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
224/// Create an active release.
225pub 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
231/// Create a pending release.
232pub 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}