Skip to main content

mvm_core/
pool.rs

1use std::fmt;
2
3use serde::{Deserialize, Serialize};
4
5use crate::tenant::tenant_pools_dir;
6
7// ============================================================================
8// Role-based VM type
9// ============================================================================
10
11/// Role for a pool's instances. Determines services, ports, drive
12/// expectations, reconcile ordering, and sleep policy.
13#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "kebab-case")]
15pub enum Role {
16    Gateway,
17    #[default]
18    Worker,
19    Builder,
20    CapabilityImessage,
21}
22
23impl fmt::Display for Role {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        match self {
26            Self::Gateway => write!(f, "gateway"),
27            Self::Worker => write!(f, "worker"),
28            Self::Builder => write!(f, "builder"),
29            Self::CapabilityImessage => write!(f, "capability-imessage"),
30        }
31    }
32}
33
34// ============================================================================
35// Minimum runtime policy
36// ============================================================================
37
38// ============================================================================
39// Pool metadata
40// ============================================================================
41
42/// Optional metadata for categorizing and tagging pools.
43/// Enables capability-based queries and policies without hardcoding types in Role enum.
44#[derive(Debug, Clone, Default, Serialize, Deserialize)]
45pub struct PoolMetadata {
46    /// Capability identifier (e.g., "openclaw", "mcp-server", "database").
47    /// Used for grouping pools by functional capability.
48    #[serde(skip_serializing_if = "Option::is_none")]
49    pub capability: Option<String>,
50
51    /// Integration types supported by this pool (e.g., ["telegram", "discord"]).
52    #[serde(default, skip_serializing_if = "Vec::is_empty")]
53    pub integration_types: Vec<String>,
54
55    /// Arbitrary key-value tags for custom categorization.
56    #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
57    pub tags: std::collections::BTreeMap<String, String>,
58}
59
60/// Per-pool runtime policy for minimum runtime enforcement and graceful lifecycle.
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct RuntimePolicy {
63    /// Minimum seconds an instance must stay Running before eligible for Warm.
64    #[serde(default = "default_min_running")]
65    pub min_running_seconds: u64,
66    /// Minimum seconds an instance must stay Warm before eligible for Sleep.
67    #[serde(default = "default_min_warm")]
68    pub min_warm_seconds: u64,
69    /// Maximum seconds to wait for guest drain ACK before forcing sleep.
70    #[serde(default = "default_drain_timeout")]
71    pub drain_timeout_seconds: u64,
72    /// Maximum seconds for graceful shutdown before SIGKILL.
73    #[serde(default = "default_graceful_shutdown")]
74    pub graceful_shutdown_seconds: u64,
75}
76
77fn default_min_running() -> u64 {
78    60
79}
80fn default_min_warm() -> u64 {
81    30
82}
83fn default_drain_timeout() -> u64 {
84    30
85}
86fn default_graceful_shutdown() -> u64 {
87    15
88}
89
90impl Default for RuntimePolicy {
91    fn default() -> Self {
92        Self {
93            min_running_seconds: default_min_running(),
94            min_warm_seconds: default_min_warm(),
95            drain_timeout_seconds: default_drain_timeout(),
96            graceful_shutdown_seconds: default_graceful_shutdown(),
97        }
98    }
99}
100
101// ============================================================================
102// Sleep policy configuration (optional per-pool override)
103// ============================================================================
104
105/// Configurable thresholds for per-pool sleep policy.
106///
107/// When set on a `DesiredPool`, overrides the system-wide default thresholds
108/// for idle detection and sleep transitions.
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct SleepPolicyConfig {
111    /// Idle seconds before Running → Warm transition (default: 300).
112    #[serde(default = "default_warm_threshold")]
113    pub warm_threshold_secs: u64,
114    /// Idle seconds before Warm → Sleeping transition (default: 900).
115    #[serde(default = "default_sleep_threshold")]
116    pub sleep_threshold_secs: u64,
117    /// CPU % below which an instance is considered idle (default: 5.0).
118    #[serde(default = "default_cpu_threshold")]
119    pub cpu_threshold: f32,
120    /// Net bytes below which an instance is considered idle (default: 1024).
121    #[serde(default = "default_net_threshold")]
122    pub net_bytes_threshold: u64,
123}
124
125fn default_warm_threshold() -> u64 {
126    300
127}
128fn default_sleep_threshold() -> u64 {
129    900
130}
131fn default_cpu_threshold() -> f32 {
132    5.0
133}
134fn default_net_threshold() -> u64 {
135    1024
136}
137
138impl Default for SleepPolicyConfig {
139    fn default() -> Self {
140        Self {
141            warm_threshold_secs: default_warm_threshold(),
142            sleep_threshold_secs: default_sleep_threshold(),
143            cpu_threshold: default_cpu_threshold(),
144            net_bytes_threshold: default_net_threshold(),
145        }
146    }
147}
148
149// ============================================================================
150// Update strategy (shared with mvmd for rollout orchestration)
151// ============================================================================
152
153/// Strategy for rolling out artifact updates to a pool's instances.
154#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
155#[serde(tag = "type", rename_all = "snake_case")]
156pub enum UpdateStrategy {
157    /// Replace instances one at a time (or in small batches).
158    Rolling(RollingUpdateStrategy),
159    /// Deploy a small number of canary instances first, then proceed.
160    Canary(CanaryStrategy),
161}
162
163impl Default for UpdateStrategy {
164    fn default() -> Self {
165        Self::Rolling(RollingUpdateStrategy::default())
166    }
167}
168
169/// Configuration for a rolling update.
170#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
171pub struct RollingUpdateStrategy {
172    /// Max instances being updated simultaneously.
173    #[serde(default = "default_max_unavailable")]
174    pub max_unavailable: u32,
175    /// Max extra instances during rollout (surge capacity).
176    #[serde(default = "default_max_surge")]
177    pub max_surge: u32,
178    /// Seconds to wait for health check after each instance starts.
179    #[serde(default = "default_health_check_timeout")]
180    pub health_check_timeout_secs: u64,
181}
182
183impl Default for RollingUpdateStrategy {
184    fn default() -> Self {
185        Self {
186            max_unavailable: default_max_unavailable(),
187            max_surge: default_max_surge(),
188            health_check_timeout_secs: default_health_check_timeout(),
189        }
190    }
191}
192
193/// Configuration for a canary deployment.
194#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
195pub struct CanaryStrategy {
196    /// Number of canary instances to deploy.
197    #[serde(default = "default_canary_count")]
198    pub canary_count: u32,
199    /// Observation window in seconds before deciding to proceed.
200    #[serde(default = "default_canary_duration")]
201    pub canary_duration_secs: u64,
202    /// Minimum health success rate to proceed (0.0–1.0).
203    #[serde(default = "default_success_threshold")]
204    pub success_threshold: f64,
205}
206
207impl Default for CanaryStrategy {
208    fn default() -> Self {
209        Self {
210            canary_count: default_canary_count(),
211            canary_duration_secs: default_canary_duration(),
212            success_threshold: default_success_threshold(),
213        }
214    }
215}
216
217fn default_max_unavailable() -> u32 {
218    1
219}
220fn default_max_surge() -> u32 {
221    1
222}
223fn default_health_check_timeout() -> u64 {
224    60
225}
226fn default_canary_count() -> u32 {
227    1
228}
229fn default_canary_duration() -> u64 {
230    300
231}
232fn default_success_threshold() -> f64 {
233    0.95
234}
235
236// ============================================================================
237// Registry artifact reference
238// ============================================================================
239
240/// Reference to pre-built artifacts in an S3-compatible registry.
241/// When present on a DesiredPool, the agent pulls artifacts from the registry
242/// instead of running a local Nix build.
243#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
244pub struct RegistryArtifact {
245    /// Template ID in the registry (matches `mvmctl template` name).
246    pub template_id: String,
247    /// Specific revision hash to pull. When None, the registry's "current"
248    /// pointer is resolved at pull time.
249    #[serde(default)]
250    pub revision: Option<String>,
251}
252
253// ============================================================================
254// Pool spec
255// ============================================================================
256
257/// A WorkerPool defines a homogeneous group of instances within a tenant.
258/// Has desired counts but NO runtime state.
259#[derive(Debug, Clone, Serialize, Deserialize)]
260pub struct PoolSpec {
261    pub pool_id: String,
262    pub tenant_id: String,
263    pub flake_ref: String,
264    /// Guest profile name. Built-in: "minimal", "python".
265    /// Users can define custom profiles in their own flake.
266    pub profile: String,
267    /// Role for all instances in this pool.
268    #[serde(default)]
269    pub role: Role,
270    pub instance_resources: InstanceResources,
271    pub desired_counts: DesiredCounts,
272    /// Minimum runtime policy for this pool's instances.
273    #[serde(default)]
274    pub runtime_policy: RuntimePolicy,
275    /// Optional metadata for capability tagging and categorization.
276    #[serde(default)]
277    pub metadata: PoolMetadata,
278    /// "baseline" | "strict"
279    #[serde(default = "default_seccomp")]
280    pub seccomp_policy: String,
281    /// "none" | "lz4" | "zstd"
282    #[serde(default = "default_compression")]
283    pub snapshot_compression: String,
284    #[serde(default)]
285    pub metadata_enabled: bool,
286    /// If true, reconcile won't auto-sleep this pool's instances.
287    #[serde(default)]
288    pub pinned: bool,
289    /// If true, reconcile won't touch this pool at all.
290    #[serde(default)]
291    pub critical: bool,
292    /// Per-integration secret scoping. When non-empty, secrets are split
293    /// into per-integration directories on the secrets drive.
294    #[serde(default)]
295    pub secret_scopes: Vec<SecretScope>,
296    /// Optional template reference for shared base image.
297    #[serde(default)]
298    pub template_id: String,
299}
300
301/// Scoped secret delivery: only give an integration the secrets it needs.
302#[derive(Debug, Clone, Serialize, Deserialize)]
303pub struct SecretScope {
304    /// Integration name (e.g. "whatsapp", "telegram").
305    pub integration: String,
306    /// Secret key names to include for this integration.
307    /// Empty means include all keys (no filtering).
308    pub keys: Vec<String>,
309}
310
311fn default_seccomp() -> String {
312    "baseline".to_string()
313}
314
315fn default_compression() -> String {
316    "none".to_string()
317}
318
319/// Resource allocation for each instance in the pool.
320#[derive(Debug, Clone, Serialize, Deserialize)]
321pub struct InstanceResources {
322    pub vcpus: u8,
323    pub mem_mib: u32,
324    #[serde(default)]
325    pub data_disk_mib: u32,
326}
327
328/// Desired instance counts by status, evaluated by the reconcile loop.
329#[derive(Debug, Clone, Default, Serialize, Deserialize)]
330pub struct DesiredCounts {
331    pub running: u32,
332    pub warm: u32,
333    pub sleeping: u32,
334}
335
336/// A completed build revision with artifact locations.
337#[derive(Debug, Clone, Serialize, Deserialize)]
338pub struct BuildRevision {
339    pub revision_hash: String,
340    pub flake_ref: String,
341    pub flake_lock_hash: String,
342    pub artifact_paths: ArtifactPaths,
343    pub built_at: String,
344}
345
346/// Artifact file sizes in bytes for build reporting and size tracking.
347#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
348pub struct ArtifactSizes {
349    #[serde(default)]
350    pub vmlinux_bytes: u64,
351    #[serde(default)]
352    pub rootfs_bytes: u64,
353    #[serde(default, skip_serializing_if = "Option::is_none")]
354    pub initrd_bytes: Option<u64>,
355    /// Nix closure size (all transitive dependencies). Optional — only
356    /// populated when `nix path-info -S` succeeds after a build.
357    #[serde(default, skip_serializing_if = "Option::is_none")]
358    pub nix_closure_bytes: Option<u64>,
359}
360
361impl ArtifactSizes {
362    /// Total size of all artifacts in bytes.
363    pub fn total_bytes(&self) -> u64 {
364        self.vmlinux_bytes + self.rootfs_bytes + self.initrd_bytes.unwrap_or(0)
365    }
366}
367
368/// Format a byte count as a human-readable string (e.g. "45.2 MiB").
369pub fn format_bytes(bytes: u64) -> String {
370    const KIB: u64 = 1024;
371    const MIB: u64 = 1024 * KIB;
372    const GIB: u64 = 1024 * MIB;
373
374    if bytes >= GIB {
375        format!("{:.1} GiB", bytes as f64 / GIB as f64)
376    } else if bytes >= MIB {
377        format!("{:.1} MiB", bytes as f64 / MIB as f64)
378    } else if bytes >= KIB {
379        format!("{:.1} KiB", bytes as f64 / KIB as f64)
380    } else {
381        format!("{} B", bytes)
382    }
383}
384
385/// Paths to build artifacts within the pool's artifact directory.
386#[derive(Debug, Clone, Serialize, Deserialize)]
387pub struct ArtifactPaths {
388    pub vmlinux: String,
389    pub rootfs: String,
390    pub fc_base_config: String,
391    /// NixOS initrd (optional — present when the flake produces one).
392    #[serde(default, skip_serializing_if = "Option::is_none")]
393    pub initrd: Option<String>,
394    /// Artifact file sizes (optional — populated by builds after Sprint 37).
395    #[serde(default, skip_serializing_if = "Option::is_none")]
396    pub sizes: Option<ArtifactSizes>,
397}
398
399// --- Filesystem paths ---
400
401pub fn pool_dir(tenant_id: &str, pool_id: &str) -> String {
402    format!("{}/{}", tenant_pools_dir(tenant_id), pool_id)
403}
404
405pub fn pool_config_path(tenant_id: &str, pool_id: &str) -> String {
406    format!("{}/pool.json", pool_dir(tenant_id, pool_id))
407}
408
409pub fn pool_artifacts_dir(tenant_id: &str, pool_id: &str) -> String {
410    format!("{}/artifacts", pool_dir(tenant_id, pool_id))
411}
412
413pub fn pool_instances_dir(tenant_id: &str, pool_id: &str) -> String {
414    format!("{}/instances", pool_dir(tenant_id, pool_id))
415}
416
417pub fn pool_snapshots_dir(tenant_id: &str, pool_id: &str) -> String {
418    format!("{}/snapshots", pool_dir(tenant_id, pool_id))
419}
420
421/// Directory for pool-level configuration data (mounted as config drive).
422pub fn pool_config_data_dir(tenant_id: &str, pool_id: &str) -> String {
423    format!("{}/config", pool_dir(tenant_id, pool_id))
424}
425
426#[cfg(test)]
427mod tests {
428    use super::*;
429
430    #[test]
431    fn test_pool_dir_path() {
432        assert_eq!(
433            pool_dir("acme", "workers"),
434            "/var/lib/mvm/tenants/acme/pools/workers"
435        );
436    }
437
438    #[test]
439    fn test_pool_config_roundtrip() {
440        let spec = PoolSpec {
441            pool_id: "workers".to_string(),
442            tenant_id: "acme".to_string(),
443            flake_ref: "github:org/repo".to_string(),
444            profile: "minimal".to_string(),
445            role: Role::Worker,
446            instance_resources: InstanceResources {
447                vcpus: 2,
448                mem_mib: 1024,
449                data_disk_mib: 2048,
450            },
451            desired_counts: DesiredCounts {
452                running: 3,
453                warm: 1,
454                sleeping: 2,
455            },
456            runtime_policy: RuntimePolicy::default(),
457            metadata: PoolMetadata::default(),
458            seccomp_policy: "baseline".to_string(),
459            snapshot_compression: "zstd".to_string(),
460            metadata_enabled: false,
461            pinned: false,
462            critical: false,
463            secret_scopes: vec![],
464            template_id: String::new(),
465        };
466
467        let json = serde_json::to_string(&spec).unwrap();
468        let parsed: PoolSpec = serde_json::from_str(&json).unwrap();
469        assert_eq!(parsed.pool_id, "workers");
470        assert_eq!(parsed.instance_resources.vcpus, 2);
471        assert_eq!(parsed.desired_counts.running, 3);
472        assert_eq!(parsed.role, Role::Worker);
473    }
474
475    #[test]
476    fn test_role_serde_roundtrip() {
477        for (role, expected) in [
478            (Role::Gateway, "\"gateway\""),
479            (Role::Worker, "\"worker\""),
480            (Role::Builder, "\"builder\""),
481            (Role::CapabilityImessage, "\"capability-imessage\""),
482        ] {
483            let json = serde_json::to_string(&role).unwrap();
484            assert_eq!(json, expected);
485            let parsed: Role = serde_json::from_str(&json).unwrap();
486            assert_eq!(parsed, role);
487        }
488    }
489
490    #[test]
491    fn test_role_display() {
492        assert_eq!(Role::Gateway.to_string(), "gateway");
493        assert_eq!(Role::Worker.to_string(), "worker");
494        assert_eq!(Role::Builder.to_string(), "builder");
495        assert_eq!(Role::CapabilityImessage.to_string(), "capability-imessage");
496    }
497
498    #[test]
499    fn test_role_default_is_worker() {
500        assert_eq!(Role::default(), Role::Worker);
501    }
502
503    #[test]
504    fn test_runtime_policy_defaults() {
505        let p = RuntimePolicy::default();
506        assert_eq!(p.min_running_seconds, 60);
507        assert_eq!(p.min_warm_seconds, 30);
508        assert_eq!(p.drain_timeout_seconds, 30);
509        assert_eq!(p.graceful_shutdown_seconds, 15);
510    }
511
512    #[test]
513    fn test_pool_spec_backward_compat() {
514        // JSON without role/runtime_policy should deserialize with defaults
515        let json = r#"{
516            "pool_id": "workers",
517            "tenant_id": "acme",
518            "flake_ref": ".",
519            "profile": "minimal",
520            "instance_resources": {"vcpus": 1, "mem_mib": 512},
521            "desired_counts": {"running": 1, "warm": 0, "sleeping": 0}
522        }"#;
523        let parsed: PoolSpec = serde_json::from_str(json).unwrap();
524        assert_eq!(parsed.role, Role::Worker);
525        assert_eq!(parsed.runtime_policy.min_running_seconds, 60);
526    }
527
528    #[test]
529    fn test_pool_config_data_dir() {
530        assert_eq!(
531            pool_config_data_dir("acme", "gateways"),
532            "/var/lib/mvm/tenants/acme/pools/gateways/config"
533        );
534    }
535
536    #[test]
537    fn test_secret_scope_serde_roundtrip() {
538        let scopes = vec![
539            SecretScope {
540                integration: "whatsapp".to_string(),
541                keys: vec![
542                    "WHATSAPP_API_KEY".to_string(),
543                    "WHATSAPP_SECRET".to_string(),
544                ],
545            },
546            SecretScope {
547                integration: "telegram".to_string(),
548                keys: vec!["TELEGRAM_BOT_TOKEN".to_string()],
549            },
550        ];
551
552        let json = serde_json::to_string(&scopes).unwrap();
553        let parsed: Vec<SecretScope> = serde_json::from_str(&json).unwrap();
554        assert_eq!(parsed.len(), 2);
555        assert_eq!(parsed[0].integration, "whatsapp");
556        assert_eq!(parsed[0].keys.len(), 2);
557        assert_eq!(parsed[1].integration, "telegram");
558    }
559
560    #[test]
561    fn test_pool_spec_backward_compat_secret_scopes() {
562        // JSON without secret_scopes should parse fine (defaults to empty vec)
563        let json = r#"{
564            "pool_id": "workers",
565            "tenant_id": "acme",
566            "flake_ref": ".",
567            "profile": "minimal",
568            "instance_resources": {"vcpus": 1, "mem_mib": 512},
569            "desired_counts": {"running": 1, "warm": 0, "sleeping": 0}
570        }"#;
571        let parsed: PoolSpec = serde_json::from_str(json).unwrap();
572        assert!(parsed.secret_scopes.is_empty());
573    }
574
575    #[test]
576    fn test_update_strategy_default_is_rolling() {
577        let strategy = UpdateStrategy::default();
578        assert!(matches!(strategy, UpdateStrategy::Rolling(_)));
579        if let UpdateStrategy::Rolling(r) = strategy {
580            assert_eq!(r.max_unavailable, 1);
581            assert_eq!(r.max_surge, 1);
582            assert_eq!(r.health_check_timeout_secs, 60);
583        }
584    }
585
586    #[test]
587    fn test_update_strategy_rolling_serde_roundtrip() {
588        let strategy = UpdateStrategy::Rolling(RollingUpdateStrategy {
589            max_unavailable: 2,
590            max_surge: 3,
591            health_check_timeout_secs: 120,
592        });
593        let json = serde_json::to_string(&strategy).unwrap();
594        let parsed: UpdateStrategy = serde_json::from_str(&json).unwrap();
595        assert_eq!(strategy, parsed);
596    }
597
598    #[test]
599    fn test_update_strategy_canary_serde_roundtrip() {
600        let strategy = UpdateStrategy::Canary(CanaryStrategy {
601            canary_count: 3,
602            canary_duration_secs: 600,
603            success_threshold: 0.99,
604        });
605        let json = serde_json::to_string(&strategy).unwrap();
606        let parsed: UpdateStrategy = serde_json::from_str(&json).unwrap();
607        assert_eq!(strategy, parsed);
608    }
609
610    #[test]
611    fn test_canary_strategy_defaults() {
612        let c = CanaryStrategy::default();
613        assert_eq!(c.canary_count, 1);
614        assert_eq!(c.canary_duration_secs, 300);
615        assert!((c.success_threshold - 0.95).abs() < 0.001);
616    }
617
618    #[test]
619    fn test_rolling_strategy_defaults() {
620        let r = RollingUpdateStrategy::default();
621        assert_eq!(r.max_unavailable, 1);
622        assert_eq!(r.max_surge, 1);
623        assert_eq!(r.health_check_timeout_secs, 60);
624    }
625
626    #[test]
627    fn test_update_strategy_tagged_json_format() {
628        // Verify the tagged enum uses "type" field
629        let rolling = UpdateStrategy::Rolling(RollingUpdateStrategy::default());
630        let json = serde_json::to_string(&rolling).unwrap();
631        assert!(json.contains(r#""type":"rolling""#));
632
633        let canary = UpdateStrategy::Canary(CanaryStrategy::default());
634        let json = serde_json::to_string(&canary).unwrap();
635        assert!(json.contains(r#""type":"canary""#));
636    }
637
638    #[test]
639    fn test_registry_artifact_serde_roundtrip() {
640        let ra = RegistryArtifact {
641            template_id: "hello".to_string(),
642            revision: Some("abc123def".to_string()),
643        };
644        let json = serde_json::to_string(&ra).unwrap();
645        let parsed: RegistryArtifact = serde_json::from_str(&json).unwrap();
646        assert_eq!(parsed.template_id, "hello");
647        assert_eq!(parsed.revision.as_deref(), Some("abc123def"));
648    }
649
650    #[test]
651    fn test_registry_artifact_no_revision() {
652        let json = r#"{"template_id": "openclaw"}"#;
653        let parsed: RegistryArtifact = serde_json::from_str(json).unwrap();
654        assert_eq!(parsed.template_id, "openclaw");
655        assert!(parsed.revision.is_none());
656    }
657
658    #[test]
659    fn test_registry_artifact_default_revision() {
660        let ra = RegistryArtifact {
661            template_id: "hello".to_string(),
662            revision: None,
663        };
664        let json = serde_json::to_string(&ra).unwrap();
665        // revision: None should be omitted or null
666        let parsed: RegistryArtifact = serde_json::from_str(&json).unwrap();
667        assert!(parsed.revision.is_none());
668    }
669
670    #[test]
671    fn test_artifact_sizes_serde_roundtrip() {
672        let sizes = ArtifactSizes {
673            vmlinux_bytes: 12_345_678,
674            rootfs_bytes: 45_678_901,
675            initrd_bytes: Some(2_345_678),
676            nix_closure_bytes: Some(100_000_000),
677        };
678        let json = serde_json::to_string(&sizes).unwrap();
679        let parsed: ArtifactSizes = serde_json::from_str(&json).unwrap();
680        assert_eq!(parsed, sizes);
681    }
682
683    #[test]
684    fn test_artifact_sizes_default() {
685        let sizes = ArtifactSizes::default();
686        assert_eq!(sizes.vmlinux_bytes, 0);
687        assert_eq!(sizes.rootfs_bytes, 0);
688        assert!(sizes.initrd_bytes.is_none());
689        assert!(sizes.nix_closure_bytes.is_none());
690    }
691
692    #[test]
693    fn test_artifact_sizes_total_bytes() {
694        let sizes = ArtifactSizes {
695            vmlinux_bytes: 100,
696            rootfs_bytes: 200,
697            initrd_bytes: Some(50),
698            nix_closure_bytes: None,
699        };
700        assert_eq!(sizes.total_bytes(), 350);
701
702        let no_initrd = ArtifactSizes {
703            vmlinux_bytes: 100,
704            rootfs_bytes: 200,
705            initrd_bytes: None,
706            nix_closure_bytes: None,
707        };
708        assert_eq!(no_initrd.total_bytes(), 300);
709    }
710
711    #[test]
712    fn test_artifact_sizes_backward_compat() {
713        // JSON without sizes field should deserialize ArtifactPaths fine
714        let json = r#"{
715            "vmlinux": "vmlinux",
716            "rootfs": "rootfs.ext4",
717            "fc_base_config": "fc-base.json"
718        }"#;
719        let parsed: ArtifactPaths = serde_json::from_str(json).unwrap();
720        assert!(parsed.sizes.is_none());
721    }
722
723    #[test]
724    fn test_artifact_paths_with_sizes() {
725        let paths = ArtifactPaths {
726            vmlinux: "vmlinux".to_string(),
727            rootfs: "rootfs.ext4".to_string(),
728            fc_base_config: "fc-base.json".to_string(),
729            initrd: None,
730            sizes: Some(ArtifactSizes {
731                vmlinux_bytes: 10_000_000,
732                rootfs_bytes: 50_000_000,
733                initrd_bytes: None,
734                nix_closure_bytes: None,
735            }),
736        };
737        let json = serde_json::to_string(&paths).unwrap();
738        let parsed: ArtifactPaths = serde_json::from_str(&json).unwrap();
739        assert!(parsed.sizes.is_some());
740        assert_eq!(parsed.sizes.unwrap().rootfs_bytes, 50_000_000);
741    }
742
743    #[test]
744    fn test_format_bytes_zero() {
745        assert_eq!(format_bytes(0), "0 B");
746    }
747
748    #[test]
749    fn test_format_bytes_bytes() {
750        assert_eq!(format_bytes(512), "512 B");
751        assert_eq!(format_bytes(1023), "1023 B");
752    }
753
754    #[test]
755    fn test_format_bytes_kib() {
756        assert_eq!(format_bytes(1024), "1.0 KiB");
757        assert_eq!(format_bytes(1536), "1.5 KiB");
758    }
759
760    #[test]
761    fn test_format_bytes_mib() {
762        assert_eq!(format_bytes(1024 * 1024), "1.0 MiB");
763        assert_eq!(format_bytes(47_400_000), "45.2 MiB");
764    }
765
766    #[test]
767    fn test_format_bytes_gib() {
768        assert_eq!(format_bytes(1024 * 1024 * 1024), "1.0 GiB");
769        assert_eq!(format_bytes(2_684_354_560), "2.5 GiB");
770    }
771}