2 files changed, 28 insertions, 3 deletions

diff --git a/src/entity_switch.rs b/src/entity_switch.rs
index 0277288..4bf77eb 100644
--- a/src/entity_switch.rs
+++ b/src/entity_switch.rs
@@ -8,10 +8,21 @@ pub enum SwitchClass {
     Switch,
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct SwitchConfig {
     pub common: EntityCommonConfig,
     pub class: SwitchClass,
+    pub publish_on_command: bool,
+}
+
+impl Default for SwitchConfig {
+    fn default() -> Self {
+        Self {
+            common: Default::default(),
+            class: Default::default(),
+            publish_on_command: true,
+        }
+    }
 }
 
 impl SwitchConfig {
@@ -40,6 +51,13 @@ impl<'a> Switch<'a> {
         })
     }
 
+    pub fn command(&self) -> Option<BinaryState> {
+        self.0.with_data(|data| {
+            let storage = data.storage.as_switch_mut();
+            storage.command.as_ref().map(|s| s.value)
+        })
+    }
+
     pub fn toggle(&mut self) -> BinaryState {
         let new_state = self.state().unwrap_or(BinaryState::Off).flip();
         self.set(new_state);
@@ -59,7 +77,7 @@ impl<'a> Switch<'a> {
     pub async fn wait(&mut self) -> BinaryState {
         loop {
             self.0.wait_command().await;
-            if let Some(state) = self.state() {
+            if let Some(state) = self.command() {
                 return state;
             }
         }
diff --git a/src/lib.rs b/src/lib.rs
index 4fae2ba..036ae12 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -219,6 +219,7 @@ pub struct SwitchState {
 pub struct SwitchStorage {
     pub state: Option<SwitchState>,
     pub command: Option<SwitchCommand>,
+    pub publish_on_command: bool,
 }
 
 #[derive(Debug)]
@@ -460,7 +461,13 @@ impl<'a> Device<'a> {
         entity_config.id = id;
         config.populate(&mut entity_config);
 
-        let entity = self.create_entity(entity_config, EntityStorage::Switch(Default::default()));
+        let entity = self.create_entity(
+            entity_config,
+            EntityStorage::Switch(SwitchStorage {
+                publish_on_command: config.publish_on_command,
+                ..Default::default()
+            }),
+        );
         Switch::new(entity)
     }