1use serde::Serialize;
2use std::collections::HashMap;
3use std::fmt;
4
5#[derive(Debug)]
6pub enum Sensor {
7 CpuUsage,
8 CpuTemperature,
9 NetRx,
10 NetTx,
11}
12
13impl Sensor {
14 pub fn as_str(&self) -> &'static str {
15 match self {
16 Sensor::CpuUsage => "cpu_usage",
17 Sensor::CpuTemperature => "cpu_temp",
18 Sensor::NetRx => "net_rx",
19 Sensor::NetTx => "net_tx",
20 }
21 }
22}
23
24#[derive(Serialize, Debug)]
25pub struct RegistrationDescriptor {
26 device: Device,
27 origin: Origin,
28 components: HashMap<&'static str, DeviceComponent>,
29 state_topic: String,
30}
31
32#[derive(Serialize, Debug)]
33pub struct Device {
34 name: String,
35 identifiers: String,
36}
37
38#[derive(Serialize, Debug)]
39pub struct Origin {
40 name: &'static str,
41 sw_version: &'static str,
42 url: &'static str,
43}
44
45#[derive(Serialize, Debug)]
46pub struct DeviceComponent {
47 name: &'static str,
48 platform: &'static str,
49 device_class: Option<&'static str>,
50 #[serde(skip_serializing_if = "Option::is_none")]
51 icon: Option<&'static str>,
52 state_class: &'static str,
53 unit_of_measurement: &'static str,
54 unique_id: String,
55 value_template: &'static str,
56 expire_after: u64,
57}
58
59impl RegistrationDescriptor {
60 pub fn new(entity: &str) -> RegistrationDescriptor {
61 let version = env!("CARGO_PKG_VERSION");
62 let package_name = env!("CARGO_PKG_NAME");
63 let url = env!("CARGO_PKG_HOMEPAGE");
64
65 RegistrationDescriptor {
66 device: Device {
67 name: entity.to_string(),
68 identifiers: entity.to_string(),
69 },
70 origin: Origin {
71 name: package_name,
72 sw_version: version,
73 url,
74 },
75 components: Default::default(),
76 state_topic: format!("mqtt-system-monitor/{entity}/state"),
77 }
78 }
79
80 pub fn add_component(&mut self, sensor: Sensor, entity: &str) {
81 self.components
82 .insert(sensor.as_str(), DeviceComponent::new(sensor, entity));
83 }
84
85 pub fn has_sensor(&self, sensor: Sensor) -> bool {
86 self.components.contains_key(sensor.as_str())
87 }
88
89 pub fn remove_sensor(&mut self, sensor: Sensor) {
90 self.components.remove(sensor.as_str());
91 }
92
93 pub fn discovery_topic(&self, prefix: &str) -> String {
94 format!("{prefix}/device/{}/config", self.device.name)
95 }
96}
97
98impl fmt::Display for RegistrationDescriptor {
99 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
100 let Ok(descriptor) = serde_json::to_string(&self) else {
101 return Err(fmt::Error);
102 };
103 write!(f, "{descriptor}")
104 }
105}
106
107impl DeviceComponent {
108 pub fn new(sensor: Sensor, entity: &str) -> DeviceComponent {
109 match sensor {
110 Sensor::CpuTemperature => Self::cpu_temperature(entity),
111 Sensor::CpuUsage => Self::cpu_usage(entity),
112 Sensor::NetRx => Self::net_rx(entity),
113 Sensor::NetTx => Self::net_tx(entity),
114 }
115 }
116 pub fn cpu_temperature(entity: &str) -> DeviceComponent {
117 DeviceComponent {
118 name: "CPU temperature",
119 platform: "sensor",
120 device_class: Some("temperature"),
121 icon: None,
122 state_class: "measurement",
123 unit_of_measurement: "°C",
124 unique_id: format!("{entity}_cpu_temp"),
125 value_template: "{{ value_json.cpu_temp }}",
126 expire_after: 60,
127 }
128 }
129
130 pub fn cpu_usage(entity: &str) -> DeviceComponent {
131 DeviceComponent {
132 name: "CPU usage",
133 platform: "sensor",
134 device_class: None,
135 state_class: "measurement",
136 icon: Some("mdi:cpu-64-bit"),
137 unit_of_measurement: "%",
138 unique_id: format!("{entity}_cpu_usage"),
139 value_template: "{{ value_json.cpu_usage }}",
140 expire_after: 60,
141 }
142 }
143
144 pub fn net_rx(entity: &str) -> DeviceComponent {
145 DeviceComponent {
146 name: "Network RX rate",
147 platform: "sensor",
148 device_class: Some("data_rate"),
149 state_class: "measurement",
150 icon: None,
151 unit_of_measurement: "KiB/s",
152 unique_id: format!("{entity}_net_rx"),
153 value_template: "{{ value_json.net_rx }}",
154 expire_after: 60,
155 }
156 }
157
158 pub fn net_tx(entity: &str) -> DeviceComponent {
159 DeviceComponent {
160 name: "Network TX rate",
161 platform: "sensor",
162 device_class: Some("data_rate"),
163 state_class: "measurement",
164 icon: None,
165 unit_of_measurement: "KiB/s",
166 unique_id: format!("{entity}_net_tx"),
167 value_template: "{{ value_json.net_tx }}",
168 expire_after: 60,
169 }
170 }
171}
172
173#[cfg(test)]
174mod tests {
175 use crate::home_assistant::{RegistrationDescriptor, Sensor};
176
177 #[test]
178 fn test_registration() -> Result<(), Box<dyn std::error::Error>> {
179 let entity = "test_entity";
180 let mut descriptor = RegistrationDescriptor::new(entity);
181
182 descriptor.add_component(Sensor::CpuUsage, entity);
183 descriptor.add_component(Sensor::CpuTemperature, entity);
184 descriptor.add_component(Sensor::NetTx, entity);
185 descriptor.add_component(Sensor::NetRx, entity);
186
187 assert_eq!(descriptor.device.name, "test_entity");
188 assert_eq!(descriptor.device.identifiers, "test_entity");
189
190 assert_eq!(
191 descriptor.state_topic,
192 "mqtt-system-monitor/test_entity/state"
193 );
194
195 for component in &descriptor.components {
196 assert_eq!(component.1.unique_id, format!("{entity}_{}", component.0));
197 assert_eq!(
198 component.1.value_template,
199 format!("{{{{ value_json.{} }}}}", component.0).as_str()
200 );
201 assert_eq!(component.1.state_class, "measurement");
202 }
203
204 let cpu_usage = descriptor
205 .components
206 .get("cpu_usage")
207 .expect("component cpu_usage not found");
208
209 assert_eq!(cpu_usage.device_class, None);
210
211 Ok(())
212 }
213}