From e6a6b2a6802bb1ff625fc0ec759d1a061e15ab86 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Wed, 3 Mar 2021 10:13:04 -0800
Subject: [PATCH] Simplify switch light (#47317)

---
 homeassistant/components/switch/light.py | 50 +++++++++---------------
 1 file changed, 19 insertions(+), 31 deletions(-)

diff --git a/homeassistant/components/switch/light.py b/homeassistant/components/switch/light.py
index 5128a49d8b7..2650bd61bfb 100644
--- a/homeassistant/components/switch/light.py
+++ b/homeassistant/components/switch/light.py
@@ -12,7 +12,7 @@ from homeassistant.const import (
     STATE_ON,
     STATE_UNAVAILABLE,
 )
-from homeassistant.core import CALLBACK_TYPE, callback
+from homeassistant.core import State, callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.event import async_track_state_change_event
@@ -37,7 +37,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
 async def async_setup_platform(
     hass: HomeAssistantType,
     config: ConfigType,
-    async_add_entities: Callable[[Sequence[Entity], bool], None],
+    async_add_entities: Callable[[Sequence[Entity]], None],
     discovery_info: Optional[DiscoveryInfoType] = None,
 ) -> None:
     """Initialize Light Switch platform."""
@@ -53,8 +53,7 @@ async def async_setup_platform(
                 config[CONF_ENTITY_ID],
                 unique_id,
             )
-        ],
-        True,
+        ]
     )
 
 
@@ -66,9 +65,7 @@ class LightSwitch(LightEntity):
         self._name = name
         self._switch_entity_id = switch_entity_id
         self._unique_id = unique_id
-        self._is_on = False
-        self._available = False
-        self._async_unsub_state_changed: Optional[CALLBACK_TYPE] = None
+        self._switch_state: Optional[State] = None
 
     @property
     def name(self) -> str:
@@ -78,12 +75,16 @@ class LightSwitch(LightEntity):
     @property
     def is_on(self) -> bool:
         """Return true if light switch is on."""
-        return self._is_on
+        assert self._switch_state is not None
+        return self._switch_state.state == STATE_ON
 
     @property
     def available(self) -> bool:
         """Return true if light switch is on."""
-        return self._available
+        return (
+            self._switch_state is not None
+            and self._switch_state.state != STATE_UNAVAILABLE
+        )
 
     @property
     def should_poll(self) -> bool:
@@ -117,33 +118,20 @@ class LightSwitch(LightEntity):
             context=self._context,
         )
 
-    async def async_update(self):
-        """Query the switch in this light switch and determine the state."""
-        switch_state = self.hass.states.get(self._switch_entity_id)
-
-        if switch_state is None:
-            self._available = False
-            return
-
-        self._is_on = switch_state.state == STATE_ON
-        self._available = switch_state.state != STATE_UNAVAILABLE
-
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""
+        assert self.hass is not None
+        self._switch_state = self.hass.states.get(self._switch_entity_id)
 
         @callback
         def async_state_changed_listener(*_: Any) -> None:
             """Handle child updates."""
-            self.async_schedule_update_ha_state(True)
+            assert self.hass is not None
+            self._switch_state = self.hass.states.get(self._switch_entity_id)
+            self.async_write_ha_state()
 
-        assert self.hass is not None
-        self._async_unsub_state_changed = async_track_state_change_event(
-            self.hass, [self._switch_entity_id], async_state_changed_listener
+        self.async_on_remove(
+            async_track_state_change_event(
+                self.hass, [self._switch_entity_id], async_state_changed_listener
+            )
         )
-
-    async def async_will_remove_from_hass(self):
-        """Handle removal from Home Assistant."""
-        if self._async_unsub_state_changed is not None:
-            self._async_unsub_state_changed()
-            self._async_unsub_state_changed = None
-            self._available = False