formatting and improved timeout handling

7 files changed, 238 insertions, 89 deletions

diff --git a/examples/button.rs b/examples/button.rs
index 4a2a228..ea8c4a9 100644
--- a/examples/button.rs
+++ b/examples/button.rs
@@ -20,7 +20,11 @@ async fn main_task(spawner: Spawner) {
         },
     );
 
-    let button = embassy_ha::create_button(&device, "button-sensor-id", embassy_ha::ButtonConfig::default());
+    let button = embassy_ha::create_button(
+        &device,
+        "button-sensor-id",
+        embassy_ha::ButtonConfig::default(),
+    );
 
     spawner.must_spawn(button_task(button));
 
diff --git a/src/entity_category.rs b/src/entity_category.rs
index 741e7dd..2560fd4 100644
--- a/src/entity_category.rs
+++ b/src/entity_category.rs
@@ -14,4 +14,3 @@ impl EntityCategory {
         }
     }
 }
-
diff --git a/src/entity_sensor.rs b/src/entity_sensor.rs
index 1a99754..1168c37 100644
--- a/src/entity_sensor.rs
+++ b/src/entity_sensor.rs
@@ -1,6 +1,4 @@
-use crate::{
-    Entity, EntityCommonConfig, EntityConfig, NumericSensorState, constants,
-};
+use crate::{Entity, EntityCommonConfig, EntityConfig, NumericSensorState, constants};
 
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
 pub enum StateClass {
@@ -90,7 +88,9 @@ impl SensorClass {
             SensorClass::Other(s) => Some(s),
             SensorClass::ApparentPower => Some(constants::HA_DEVICE_CLASS_SENSOR_APPARENT_POWER),
             SensorClass::Aqi => Some(constants::HA_DEVICE_CLASS_SENSOR_AQI),
-            SensorClass::AtmosphericPressure => Some(constants::HA_DEVICE_CLASS_SENSOR_ATMOSPHERIC_PRESSURE),
+            SensorClass::AtmosphericPressure => {
+                Some(constants::HA_DEVICE_CLASS_SENSOR_ATMOSPHERIC_PRESSURE)
+            }
             SensorClass::Battery => Some(constants::HA_DEVICE_CLASS_SENSOR_BATTERY),
             SensorClass::CarbonDioxide => Some(constants::HA_DEVICE_CLASS_SENSOR_CARBON_DIOXIDE),
             SensorClass::CarbonMonoxide => Some(constants::HA_DEVICE_CLASS_SENSOR_CARBON_MONOXIDE),
@@ -110,8 +110,12 @@ impl SensorClass {
             SensorClass::Irradiance => Some(constants::HA_DEVICE_CLASS_SENSOR_IRRADIANCE),
             SensorClass::Moisture => Some(constants::HA_DEVICE_CLASS_SENSOR_MOISTURE),
             SensorClass::Monetary => Some(constants::HA_DEVICE_CLASS_SENSOR_MONETARY),
-            SensorClass::NitrogenDioxide => Some(constants::HA_DEVICE_CLASS_SENSOR_NITROGEN_DIOXIDE),
-            SensorClass::NitrogenMonoxide => Some(constants::HA_DEVICE_CLASS_SENSOR_NITROGEN_MONOXIDE),
+            SensorClass::NitrogenDioxide => {
+                Some(constants::HA_DEVICE_CLASS_SENSOR_NITROGEN_DIOXIDE)
+            }
+            SensorClass::NitrogenMonoxide => {
+                Some(constants::HA_DEVICE_CLASS_SENSOR_NITROGEN_MONOXIDE)
+            }
             SensorClass::NitrousOxide => Some(constants::HA_DEVICE_CLASS_SENSOR_NITROUS_OXIDE),
             SensorClass::Ozone => Some(constants::HA_DEVICE_CLASS_SENSOR_OZONE),
             SensorClass::Ph => Some(constants::HA_DEVICE_CLASS_SENSOR_PH),
@@ -121,7 +125,9 @@ impl SensorClass {
             SensorClass::PowerFactor => Some(constants::HA_DEVICE_CLASS_SENSOR_POWER_FACTOR),
             SensorClass::Power => Some(constants::HA_DEVICE_CLASS_SENSOR_POWER),
             SensorClass::Precipitation => Some(constants::HA_DEVICE_CLASS_SENSOR_PRECIPITATION),
-            SensorClass::PrecipitationIntensity => Some(constants::HA_DEVICE_CLASS_SENSOR_PRECIPITATION_INTENSITY),
+            SensorClass::PrecipitationIntensity => {
+                Some(constants::HA_DEVICE_CLASS_SENSOR_PRECIPITATION_INTENSITY)
+            }
             SensorClass::Pressure => Some(constants::HA_DEVICE_CLASS_SENSOR_PRESSURE),
             SensorClass::ReactivePower => Some(constants::HA_DEVICE_CLASS_SENSOR_REACTIVE_POWER),
             SensorClass::SignalStrength => Some(constants::HA_DEVICE_CLASS_SENSOR_SIGNAL_STRENGTH),
@@ -130,8 +136,12 @@ impl SensorClass {
             SensorClass::SulphurDioxide => Some(constants::HA_DEVICE_CLASS_SENSOR_SULPHUR_DIOXIDE),
             SensorClass::Temperature => Some(constants::HA_DEVICE_CLASS_SENSOR_TEMPERATURE),
             SensorClass::Timestamp => Some(constants::HA_DEVICE_CLASS_SENSOR_TIMESTAMP),
-            SensorClass::VolatileOrganicCompounds => Some(constants::HA_DEVICE_CLASS_SENSOR_VOLATILE_ORGANIC_COMPOUNDS),
-            SensorClass::VolatileOrganicCompoundsParts => Some(constants::HA_DEVICE_CLASS_SENSOR_VOLATILE_ORGANIC_COMPOUNDS_PARTS),
+            SensorClass::VolatileOrganicCompounds => {
+                Some(constants::HA_DEVICE_CLASS_SENSOR_VOLATILE_ORGANIC_COMPOUNDS)
+            }
+            SensorClass::VolatileOrganicCompoundsParts => {
+                Some(constants::HA_DEVICE_CLASS_SENSOR_VOLATILE_ORGANIC_COMPOUNDS_PARTS)
+            }
             SensorClass::Voltage => Some(constants::HA_DEVICE_CLASS_SENSOR_VOLTAGE),
             SensorClass::Volume => Some(constants::HA_DEVICE_CLASS_SENSOR_VOLUME),
             SensorClass::VolumeFlowRate => Some(constants::HA_DEVICE_CLASS_SENSOR_VOLUME_FLOW_RATE),
diff --git a/src/entity_switch.rs b/src/entity_switch.rs
index 1cb3647..299d299 100644
--- a/src/entity_switch.rs
+++ b/src/entity_switch.rs
@@ -1,4 +1,6 @@
-use crate::{BinaryState, Entity, EntityCommonConfig, EntityConfig, SwitchCommand, SwitchState, constants};
+use crate::{
+    BinaryState, Entity, EntityCommonConfig, EntityConfig, SwitchCommand, SwitchState, constants,
+};
 
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
 pub enum SwitchClass {
diff --git a/src/lib.rs b/src/lib.rs
index c1b5191..672fde9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,10 +1,14 @@
 #![no_std]
 
-use core::{cell::RefCell, net::{Ipv4Addr, SocketAddrV4}, task::Waker};
+use core::{
+    cell::RefCell,
+    net::{Ipv4Addr, SocketAddrV4},
+    task::Waker,
+};
 
 use embassy_net::tcp::TcpSocket;
 use embassy_sync::waitqueue::AtomicWaker;
-use embassy_time::Timer;
+use embassy_time::{Duration, Timer};
 use heapless::{
     Vec, VecView,
     string::{String, StringView},
@@ -432,7 +436,11 @@ pub fn new<'a>(resources: &'a mut DeviceResources, config: DeviceConfig) -> Devi
     }
 }
 
-fn create_entity<'a>(device: &Device<'a>, config: EntityConfig, storage: EntityStorage) -> Entity<'a> {
+fn create_entity<'a>(
+    device: &Device<'a>,
+    config: EntityConfig,
+    storage: EntityStorage,
+) -> Entity<'a> {
     let index = 'outer: {
         for idx in 0..device.entities.len() {
             if device.entities[idx].borrow().is_none() {
@@ -457,7 +465,11 @@ fn create_entity<'a>(device: &Device<'a>, config: EntityConfig, storage: EntityS
     }
 }
 
-pub fn create_sensor<'a>(device: &Device<'a>, id: &'static str, config: SensorConfig) -> Sensor<'a> {
+pub fn create_sensor<'a>(
+    device: &Device<'a>,
+    id: &'static str,
+    config: SensorConfig,
+) -> Sensor<'a> {
     let mut entity_config = EntityConfig::default();
     entity_config.id = id;
     config.populate(&mut entity_config);
@@ -470,16 +482,28 @@ pub fn create_sensor<'a>(device: &Device<'a>, id: &'static str, config: SensorCo
     Sensor::new(entity)
 }
 
-pub fn create_button<'a>(device: &Device<'a>, id: &'static str, config: ButtonConfig) -> Button<'a> {
+pub fn create_button<'a>(
+    device: &Device<'a>,
+    id: &'static str,
+    config: ButtonConfig,
+) -> Button<'a> {
     let mut entity_config = EntityConfig::default();
     entity_config.id = id;
     config.populate(&mut entity_config);
 
-    let entity = create_entity(device, entity_config, EntityStorage::Button(Default::default()));
+    let entity = create_entity(
+        device,
+        entity_config,
+        EntityStorage::Button(Default::default()),
+    );
     Button::new(entity)
 }
 
-pub fn create_number<'a>(device: &Device<'a>, id: &'static str, config: NumberConfig) -> Number<'a> {
+pub fn create_number<'a>(
+    device: &Device<'a>,
+    id: &'static str,
+    config: NumberConfig,
+) -> Number<'a> {
     let mut entity_config = EntityConfig::default();
     entity_config.id = id;
     config.populate(&mut entity_config);
@@ -495,7 +519,11 @@ pub fn create_number<'a>(device: &Device<'a>, id: &'static str, config: NumberCo
     Number::new(entity)
 }
 
-pub fn create_switch<'a>(device: &Device<'a>, id: &'static str, config: SwitchConfig) -> Switch<'a> {
+pub fn create_switch<'a>(
+    device: &Device<'a>,
+    id: &'static str,
+    config: SwitchConfig,
+) -> Switch<'a> {
     let mut entity_config = EntityConfig::default();
     entity_config.id = id;
     config.populate(&mut entity_config);
@@ -531,6 +559,8 @@ pub fn create_binary_sensor<'a>(
 pub async fn run<T: Transport>(device: &mut Device<'_>, transport: &mut T) -> Result<(), Error> {
     use core::fmt::Write;
 
+    const MQTT_TIMEOUT: Duration = Duration::from_secs(30);
+
     device.availability_topic_buffer.clear();
     write!(
         device.availability_topic_buffer,
@@ -549,15 +579,24 @@ pub async fn run<T: Transport>(device: &mut Device<'_>, transport: &mut T) -> Re
         will_retain: true,
         ..Default::default()
     };
-    if let Err(err) = client
-        .connect_with(device.config.device_id, connect_params)
-        .await
+    match embassy_time::with_timeout(
+        MQTT_TIMEOUT,
+        client.connect_with(device.config.device_id, connect_params),
+    )
+    .await
     {
-        crate::log::error!(
-            "mqtt connect failed with: {:?}",
-            crate::log::Debug2Format(&err)
-        );
-        return Err(Error::new("mqtt connection failed"));
+        Ok(Ok(())) => {}
+        Ok(Err(err)) => {
+            crate::log::error!(
+                "mqtt connect failed with: {:?}",
+                crate::log::Debug2Format(&err)
+            );
+            return Err(Error::new("mqtt connection failed"));
+        }
+        Err(_) => {
+            crate::log::error!("mqtt connect timed out");
+            return Err(Error::new("mqtt connect timed out"));
+        }
     }
 
     crate::log::debug!("sending discover messages");
@@ -635,7 +674,8 @@ pub async fn run<T: Transport>(device: &mut Device<'_>, transport: &mut T) -> Re
                 discovery
             );
 
-            device.discovery_buffer
+            device
+                .discovery_buffer
                 .resize(device.discovery_buffer.capacity(), 0)
                 .unwrap();
             let n = serde_json_core::to_slice(&discovery, &mut device.discovery_buffer)
@@ -645,48 +685,73 @@ pub async fn run<T: Transport>(device: &mut Device<'_>, transport: &mut T) -> Re
 
         let discovery_topic = device.discovery_topic_buffer.as_str();
         crate::log::debug!("sending discovery to topic '{}'", discovery_topic);
-        if let Err(err) = client
-            .publish(discovery_topic, &device.discovery_buffer)
-            .await
+        match embassy_time::with_timeout(
+            MQTT_TIMEOUT,
+            client.publish(discovery_topic, &device.discovery_buffer),
+        )
+        .await
         {
-            crate::log::error!(
-                "mqtt discovery publish failed with: {:?}",
-                crate::log::Debug2Format(&err)
-            );
-            return Err(Error::new("mqtt discovery publish failed"));
+            Ok(Ok(_)) => {}
+            Ok(Err(err)) => {
+                crate::log::error!(
+                    "mqtt discovery publish failed with: {:?}",
+                    crate::log::Debug2Format(&err)
+                );
+                return Err(Error::new("mqtt discovery publish failed"));
+            }
+            Err(_) => {
+                crate::log::error!("mqtt discovery publish timed out");
+                return Err(Error::new("mqtt discovery publish timed out"));
+            }
         }
 
         let command_topic = device.command_topic_buffer.as_str();
         crate::log::debug!("subscribing to command topic '{}'", command_topic);
-        if let Err(err) = client.subscribe(command_topic).await {
-            crate::log::error!(
-                "mqtt subscribe to '{}' failed with: {:?}",
-                command_topic,
-                crate::log::Debug2Format(&err)
-            );
-            return Err(Error::new(
-                "mqtt subscription to entity command topic failed",
-            ));
+        match embassy_time::with_timeout(MQTT_TIMEOUT, client.subscribe(command_topic)).await {
+            Ok(Ok(_)) => {}
+            Ok(Err(err)) => {
+                crate::log::error!(
+                    "mqtt subscribe to '{}' failed with: {:?}",
+                    command_topic,
+                    crate::log::Debug2Format(&err)
+                );
+                return Err(Error::new(
+                    "mqtt subscription to entity command topic failed",
+                ));
+            }
+            Err(_) => {
+                crate::log::error!("mqtt subscribe to '{}' timed out", command_topic);
+                return Err(Error::new("mqtt subscribe timed out"));
+            }
         }
     }
 
-        if let Err(err) = client
-            .publish_with(
-                availability_topic,
-                AVAILABLE_PAYLOAD.as_bytes(),
-                embedded_mqtt::PublishParams {
-                    retain: true,
-                    ..Default::default()
-                },
-            )
-            .await
-        {
+    match embassy_time::with_timeout(
+        MQTT_TIMEOUT,
+        client.publish_with(
+            availability_topic,
+            AVAILABLE_PAYLOAD.as_bytes(),
+            embedded_mqtt::PublishParams {
+                retain: true,
+                ..Default::default()
+            },
+        ),
+    )
+    .await
+    {
+        Ok(Ok(_)) => {}
+        Ok(Err(err)) => {
             crate::log::error!(
                 "mqtt availability publish failed with: {:?}",
                 crate::log::Debug2Format(&err)
             );
             return Err(Error::new("mqtt availability publish failed"));
         }
+        Err(_) => {
+            crate::log::error!("mqtt availability publish timed out");
+            return Err(Error::new("mqtt availability publish timed out"));
+        }
+    }
 
     'outer_loop: loop {
         use core::fmt::Write;
@@ -749,7 +814,14 @@ pub async fn run<T: Transport>(device: &mut Device<'_>, transport: &mut T) -> Re
             }
 
             let state_topic = device.state_topic_buffer.as_str();
-            if let Err(err) = client.publish(state_topic, device.publish_buffer).await {
+            match embassy_time::with_timeout(
+                MQTT_TIMEOUT,
+                client.publish(state_topic, device.publish_buffer),
+            )
+            .await
+            {
+                Ok(Ok(_)) => {}
+                Ok(Err(err)) => {
                     crate::log::error!(
                         "mqtt state publish on topic '{}' failed with: {:?}",
                         state_topic,
@@ -757,12 +829,22 @@ pub async fn run<T: Transport>(device: &mut Device<'_>, transport: &mut T) -> Re
                     );
                     return Err(Error::new("mqtt publish failed"));
                 }
+                Err(_) => {
+                    crate::log::error!("mqtt state publish on topic '{}' timed out", state_topic);
+                    return Err(Error::new("mqtt publish timed out"));
+                }
             }
+        }
 
         let receive = client.receive();
         let waker = wait_on_atomic_waker(device.waker);
-        let publish = match embassy_futures::select::select(receive, waker).await {
-            embassy_futures::select::Either::First(packet) => match packet {
+        let publish = match embassy_time::with_timeout(
+            MQTT_TIMEOUT,
+            embassy_futures::select::select(receive, waker),
+        )
+        .await
+        {
+            Ok(embassy_futures::select::Either::First(packet)) => match packet {
                 Ok(embedded_mqtt::Packet::Publish(publish)) => publish,
                 Err(err) => {
                     crate::log::error!(
@@ -773,7 +855,11 @@ pub async fn run<T: Transport>(device: &mut Device<'_>, transport: &mut T) -> Re
                 }
                 _ => continue,
             },
-            embassy_futures::select::Either::Second(_) => continue,
+            Ok(embassy_futures::select::Either::Second(_)) => continue,
+            Err(_) => {
+                crate::log::error!("mqtt receive timed out");
+                return Err(Error::new("mqtt receive timed out"));
+            }
         };
 
         let entity = 'entity_search_block: {
@@ -816,12 +902,24 @@ pub async fn run<T: Transport>(device: &mut Device<'_>, transport: &mut T) -> Re
         );
 
         let data_len = publish.data_len;
-        if let Err(err) = client.receive_data(&mut read_buffer[..data_len]).await {
-            crate::log::error!(
-                "mqtt receive data failed with: {:?}",
-                crate::log::Debug2Format(&err)
-            );
-            return Err(Error::new("mqtt receive data failed"));
+        match embassy_time::with_timeout(
+            MQTT_TIMEOUT,
+            client.receive_data(&mut read_buffer[..data_len]),
+        )
+        .await
+        {
+            Ok(Ok(())) => {}
+            Ok(Err(err)) => {
+                crate::log::error!(
+                    "mqtt receive data failed with: {:?}",
+                    crate::log::Debug2Format(&err)
+                );
+                return Err(Error::new("mqtt receive data failed"));
+            }
+            Err(_) => {
+                crate::log::error!("mqtt receive data timed out");
+                return Err(Error::new("mqtt receive data timed out"));
+            }
         }
 
         let command = match str::from_utf8(&read_buffer[..data_len]) {
@@ -990,13 +1088,21 @@ pub async fn connect_and_run(
         let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
         socket.set_timeout(Some(embassy_time::Duration::from_secs(10)));
 
-        if let Err(err) = socket.connect(addr).await {
-            crate::log::error!(
-                "TCP connect to {} failed with: {:?}",
-                addr,
-                crate::log::Debug2Format(&err)
-            );
-            continue;
+        let connect_fut = embassy_time::with_timeout(Duration::from_secs(10), socket.connect(addr));
+        match connect_fut.await {
+            Ok(Err(err)) => {
+                crate::log::error!(
+                    "TCP connect to {} failed with: {:?}",
+                    addr,
+                    crate::log::Debug2Format(&err)
+                );
+                continue;
+            }
+            Err(_) => {
+                crate::log::error!("TCP connect to {} timed out", addr);
+                continue;
+            }
+            _ => {}
         }
 
         socket.set_timeout(None);
diff --git a/src/log.rs b/src/log.rs
index 1f7bebd..d25d210 100644
--- a/src/log.rs
+++ b/src/log.rs
@@ -112,4 +112,4 @@ macro_rules! error {
 }
 
 // Re-export the macros at the module level for easier use
-pub use crate::{trace, debug, info, warn, error};
+pub use crate::{debug, error, info, trace, warn};
diff --git a/src/unit.rs b/src/unit.rs
index e61c867..90eecb4 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -365,7 +365,9 @@ impl NumberUnit {
             NumberUnit::MeterPerSecond => crate::constants::HA_UNIT_SPEED_METER_PER_SECOND,
             NumberUnit::MilePerHour => crate::constants::HA_UNIT_SPEED_MILE_PER_HOUR,
             NumberUnit::MillimeterPerDay => crate::constants::HA_UNIT_SPEED_MILLIMETER_PER_DAY,
-            NumberUnit::MillimeterPerSecond => crate::constants::HA_UNIT_SPEED_MILLIMETER_PER_SECOND,
+            NumberUnit::MillimeterPerSecond => {
+                crate::constants::HA_UNIT_SPEED_MILLIMETER_PER_SECOND
+            }
             // Distance
             NumberUnit::Kilometer => crate::constants::HA_UNIT_DISTANCE_KILOMETER,
             NumberUnit::Meter => crate::constants::HA_UNIT_DISTANCE_METER,
@@ -397,12 +399,24 @@ impl NumberUnit {
             NumberUnit::MegaBitPerSecond => crate::constants::HA_UNIT_DATA_RATE_MEGABIT_PER_SECOND,
             NumberUnit::GigaBitPerSecond => crate::constants::HA_UNIT_DATA_RATE_GIGABIT_PER_SECOND,
             NumberUnit::BytePerSecond => crate::constants::HA_UNIT_DATA_RATE_BYTE_PER_SECOND,
-            NumberUnit::KiloBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_KILOBYTE_PER_SECOND,
-            NumberUnit::MegaBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_MEGABYTE_PER_SECOND,
-            NumberUnit::GigaBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_GIGABYTE_PER_SECOND,
-            NumberUnit::KibiBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_KIBIBYTE_PER_SECOND,
-            NumberUnit::MebiBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_MEBIBYTE_PER_SECOND,
-            NumberUnit::GibiBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_GIBIBYTE_PER_SECOND,
+            NumberUnit::KiloBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_KILOBYTE_PER_SECOND
+            }
+            NumberUnit::MegaBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_MEGABYTE_PER_SECOND
+            }
+            NumberUnit::GigaBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_GIGABYTE_PER_SECOND
+            }
+            NumberUnit::KibiBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_KIBIBYTE_PER_SECOND
+            }
+            NumberUnit::MebiBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_MEBIBYTE_PER_SECOND
+            }
+            NumberUnit::GibiBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_GIBIBYTE_PER_SECOND
+            }
             // Weight
             NumberUnit::Kilogram => crate::constants::HA_UNIT_WEIGHT_KILOGRAM,
             NumberUnit::Gram => crate::constants::HA_UNIT_WEIGHT_GRAM,
@@ -631,7 +645,9 @@ impl Unit {
             Unit::NumberMeterPerSecond => crate::constants::HA_UNIT_SPEED_METER_PER_SECOND,
             Unit::NumberMilePerHour => crate::constants::HA_UNIT_SPEED_MILE_PER_HOUR,
             Unit::NumberMillimeterPerDay => crate::constants::HA_UNIT_SPEED_MILLIMETER_PER_DAY,
-            Unit::NumberMillimeterPerSecond => crate::constants::HA_UNIT_SPEED_MILLIMETER_PER_SECOND,
+            Unit::NumberMillimeterPerSecond => {
+                crate::constants::HA_UNIT_SPEED_MILLIMETER_PER_SECOND
+            }
             Unit::NumberKilometer => crate::constants::HA_UNIT_DISTANCE_KILOMETER,
             Unit::NumberMeter => crate::constants::HA_UNIT_DISTANCE_METER,
             Unit::NumberCentimeter => crate::constants::HA_UNIT_DISTANCE_CENTIMETER,
@@ -658,12 +674,24 @@ impl Unit {
             Unit::NumberMegaBitPerSecond => crate::constants::HA_UNIT_DATA_RATE_MEGABIT_PER_SECOND,
             Unit::NumberGigaBitPerSecond => crate::constants::HA_UNIT_DATA_RATE_GIGABIT_PER_SECOND,
             Unit::NumberBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_BYTE_PER_SECOND,
-            Unit::NumberKiloBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_KILOBYTE_PER_SECOND,
-            Unit::NumberMegaBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_MEGABYTE_PER_SECOND,
-            Unit::NumberGigaBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_GIGABYTE_PER_SECOND,
-            Unit::NumberKibiBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_KIBIBYTE_PER_SECOND,
-            Unit::NumberMebiBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_MEBIBYTE_PER_SECOND,
-            Unit::NumberGibiBytePerSecond => crate::constants::HA_UNIT_DATA_RATE_GIBIBYTE_PER_SECOND,
+            Unit::NumberKiloBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_KILOBYTE_PER_SECOND
+            }
+            Unit::NumberMegaBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_MEGABYTE_PER_SECOND
+            }
+            Unit::NumberGigaBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_GIGABYTE_PER_SECOND
+            }
+            Unit::NumberKibiBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_KIBIBYTE_PER_SECOND
+            }
+            Unit::NumberMebiBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_MEBIBYTE_PER_SECOND
+            }
+            Unit::NumberGibiBytePerSecond => {
+                crate::constants::HA_UNIT_DATA_RATE_GIBIBYTE_PER_SECOND
+            }
             Unit::NumberKilogram => crate::constants::HA_UNIT_WEIGHT_KILOGRAM,
             Unit::NumberGram => crate::constants::HA_UNIT_WEIGHT_GRAM,
             Unit::NumberMilligram => crate::constants::HA_UNIT_WEIGHT_MILLIGRAM,