zombienet_configuration/
observability.rs

1use serde::{Deserialize, Serialize};
2
3use crate::shared::types::Port;
4
5const DEFAULT_PROMETHEUS_IMAGE: &str = "prom/prometheus:latest";
6const DEFAULT_GRAFANA_IMAGE: &str = "grafana/grafana:latest";
7
8/// Configuration for the observability stack (Prometheus + Grafana)
9///
10/// When enabled, Docker/Podman containers are spawned after the network is up,
11/// auto-configured to scrape all nodes' Prometheus metrics endpoints
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13pub struct ObservabilityConfig {
14    /// Whether the observability stack is enabled
15    #[serde(default)]
16    enabled: bool,
17    /// Host port to expose Prometheus on. If `None`, a random available port is used
18    #[serde(default)]
19    prometheus_port: Option<Port>,
20    /// Host port to expose Grafana on. If `None`, a random available port is used
21    #[serde(default)]
22    grafana_port: Option<Port>,
23    /// Docker image for Prometheus
24    #[serde(default = "default_prometheus_image")]
25    prometheus_image: String,
26    /// Docker image for Grafana
27    #[serde(default = "default_grafana_image")]
28    grafana_image: String,
29}
30
31fn default_prometheus_image() -> String {
32    DEFAULT_PROMETHEUS_IMAGE.to_string()
33}
34
35fn default_grafana_image() -> String {
36    DEFAULT_GRAFANA_IMAGE.to_string()
37}
38
39impl Default for ObservabilityConfig {
40    fn default() -> Self {
41        Self {
42            enabled: false,
43            prometheus_port: None,
44            grafana_port: None,
45            prometheus_image: default_prometheus_image(),
46            grafana_image: default_grafana_image(),
47        }
48    }
49}
50
51impl ObservabilityConfig {
52    pub fn enabled(&self) -> bool {
53        self.enabled
54    }
55
56    pub fn prometheus_port(&self) -> Option<Port> {
57        self.prometheus_port
58    }
59
60    pub fn grafana_port(&self) -> Option<Port> {
61        self.grafana_port
62    }
63
64    pub fn prometheus_image(&self) -> &str {
65        &self.prometheus_image
66    }
67
68    pub fn grafana_image(&self) -> &str {
69        &self.grafana_image
70    }
71}
72
73/// Builder for [`ObservabilityConfig`]
74#[derive(Default)]
75pub struct ObservabilityConfigBuilder {
76    config: ObservabilityConfig,
77}
78
79impl ObservabilityConfigBuilder {
80    pub fn new() -> Self {
81        Self::default()
82    }
83
84    /// Enable or disable the observability stack
85    pub fn with_enabled(mut self, enabled: bool) -> Self {
86        self.config.enabled = enabled;
87        self
88    }
89
90    /// Set the host port for Prometheus
91    pub fn with_prometheus_port(mut self, port: Port) -> Self {
92        self.config.prometheus_port = Some(port);
93        self
94    }
95
96    /// Set the host port for Grafana
97    pub fn with_grafana_port(mut self, port: Port) -> Self {
98        self.config.grafana_port = Some(port);
99        self
100    }
101
102    /// Set a custom Prometheus Docker image
103    pub fn with_prometheus_image(mut self, image: impl Into<String>) -> Self {
104        self.config.prometheus_image = image.into();
105        self
106    }
107
108    /// Set a custom Grafana Docker image
109    pub fn with_grafana_image(mut self, image: impl Into<String>) -> Self {
110        self.config.grafana_image = image.into();
111        self
112    }
113
114    pub fn build(self) -> ObservabilityConfig {
115        self.config
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn default_config_is_disabled() {
125        let config = ObservabilityConfig::default();
126        assert!(!config.enabled());
127        assert_eq!(config.prometheus_port(), None);
128        assert_eq!(config.grafana_port(), None);
129        assert_eq!(config.prometheus_image(), "prom/prometheus:latest");
130        assert_eq!(config.grafana_image(), "grafana/grafana:latest");
131    }
132
133    #[test]
134    fn builder_defaults_are_disabled() {
135        let config = ObservabilityConfigBuilder::new().build();
136        assert!(!config.enabled());
137        assert_eq!(config.prometheus_port(), None);
138        assert_eq!(config.grafana_port(), None);
139    }
140
141    #[test]
142    fn builder_with_all_fields() {
143        let config = ObservabilityConfigBuilder::new()
144            .with_enabled(true)
145            .with_prometheus_port(9090)
146            .with_grafana_port(3000)
147            .with_prometheus_image("prom/prometheus:v2.50.0")
148            .with_grafana_image("grafana/grafana:10.0.0")
149            .build();
150
151        assert!(config.enabled());
152        assert_eq!(config.prometheus_port(), Some(9090));
153        assert_eq!(config.grafana_port(), Some(3000));
154        assert_eq!(config.prometheus_image(), "prom/prometheus:v2.50.0");
155        assert_eq!(config.grafana_image(), "grafana/grafana:10.0.0");
156    }
157
158    #[test]
159    fn toml_round_trip() {
160        let config = ObservabilityConfigBuilder::new()
161            .with_enabled(true)
162            .with_prometheus_port(9090)
163            .with_grafana_port(3000)
164            .build();
165
166        let toml_str = toml::to_string(&config).unwrap();
167        let deserialized: ObservabilityConfig = toml::from_str(&toml_str).unwrap();
168        assert_eq!(config, deserialized);
169    }
170
171    #[test]
172    fn deserialize_from_toml_string() {
173        let toml_str = r#"
174            enabled = true
175            prometheus_port = 9090
176            grafana_port = 3000
177            prometheus_image = "prom/prometheus:v2.50.0"
178        "#;
179
180        let config: ObservabilityConfig = toml::from_str(toml_str).unwrap();
181        assert!(config.enabled());
182        assert_eq!(config.prometheus_port(), Some(9090));
183        assert_eq!(config.grafana_port(), Some(3000));
184        assert_eq!(config.prometheus_image(), "prom/prometheus:v2.50.0");
185        assert_eq!(config.grafana_image(), "grafana/grafana:latest");
186    }
187
188    #[test]
189    fn deserialize_empty_toml_defaults_to_disabled() {
190        let config: ObservabilityConfig = toml::from_str("").unwrap();
191        assert!(!config.enabled());
192    }
193}