From f70f524c969de4c10d4db15413cbf2000dc0674d Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Tue, 29 Aug 2023 02:15:15 -0400 Subject: [PATCH 1/7] Add zwave_js controller status sensor --- homeassistant/components/zwave_js/__init__.py | 9 +- homeassistant/components/zwave_js/api.py | 42 +++++++++ homeassistant/components/zwave_js/sensor.py | 93 ++++++++++++++++++- .../zwave_js/fixtures/controller_state.json | 3 +- tests/components/zwave_js/test_api.py | 46 +++++++++ tests/components/zwave_js/test_sensor.py | 45 +++++++++ 6 files changed, 232 insertions(+), 6 deletions(-) diff --git a/homeassistant/components/zwave_js/__init__.py b/homeassistant/components/zwave_js/__init__.py index 2e6ff4f0b34fd..a49363c0f3887 100644 --- a/homeassistant/components/zwave_js/__init__.py +++ b/homeassistant/components/zwave_js/__init__.py @@ -347,8 +347,13 @@ async def async_on_node_added(self, node: ZwaveNode) -> None: ) ) - # No need for a ping button or node status sensor for controller nodes - if not node.is_controller_node: + if node.is_controller_node: + # Create a controller status sensor for each device + async_dispatcher_send( + self.hass, + f"{DOMAIN}_{self.config_entry.entry_id}_add_controller_status_sensor", + ) + else: # Create a node status sensor for each device async_dispatcher_send( self.hass, diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py index 6781ccacdc7ed..7d1a9a40c3db7 100644 --- a/homeassistant/components/zwave_js/api.py +++ b/homeassistant/components/zwave_js/api.py @@ -391,6 +391,7 @@ def node_status(node: Node) -> dict[str, Any]: def async_register_api(hass: HomeAssistant) -> None: """Register all of our api endpoints.""" websocket_api.async_register_command(hass, websocket_network_status) + websocket_api.async_register_command(hass, websocket_subscribe_controller_status) websocket_api.async_register_command(hass, websocket_subscribe_node_status) websocket_api.async_register_command(hass, websocket_node_status) websocket_api.async_register_command(hass, websocket_node_metadata) @@ -523,6 +524,47 @@ async def websocket_network_status( ) +@websocket_api.websocket_command( + { + vol.Required(TYPE): "zwave_js/subscribe_controller_status", + vol.Required(DEVICE_ID): str, + } +) +@websocket_api.async_response +@async_get_node +async def websocket_subscribe_controller_status( + hass: HomeAssistant, + connection: ActiveConnection, + msg: dict[str, Any], + node: Node, +) -> None: + """Subscribe to controller status update events of a Z-Wave JS controller.""" + driver = node.client.driver + assert driver + controller = driver.controller + + @callback + def forward_event(event: dict) -> None: + """Forward the event.""" + connection.send_message( + websocket_api.event_message( + msg[ID], + {"event": event["event"], "status": controller.status}, + ) + ) + + @callback + def async_cleanup() -> None: + """Remove signal listeners.""" + for unsub in unsubs: + unsub() + + connection.subscriptions[msg[ID]] = async_cleanup + msg[DATA_UNSUBSCRIBE] = unsubs = [controller.on("status changed", forward_event)] + + connection.send_result(msg[ID], controller.status) + + @websocket_api.websocket_command( { vol.Required(TYPE): "zwave_js/subscribe_node_status", diff --git a/homeassistant/components/zwave_js/sensor.py b/homeassistant/components/zwave_js/sensor.py index 468d8f0cbda17..cdc4deb5329d3 100644 --- a/homeassistant/components/zwave_js/sensor.py +++ b/homeassistant/components/zwave_js/sensor.py @@ -6,7 +6,7 @@ import voluptuous as vol from zwave_js_server.client import Client as ZwaveClient -from zwave_js_server.const import CommandClass, NodeStatus +from zwave_js_server.const import CommandClass, ControllerStatus, NodeStatus from zwave_js_server.const.command_class.meter import ( RESET_METER_OPTION_TARGET_VALUE, RESET_METER_OPTION_TYPE, @@ -91,7 +91,13 @@ PARALLEL_UPDATES = 0 -STATUS_ICON: dict[NodeStatus, str] = { +CONTROLLER_STATUS_ICON: dict[ControllerStatus, str] = { + ControllerStatus.READY: "mdi:check", + ControllerStatus.UNRESPONSIVE: "mdi:bell-off", + ControllerStatus.JAMMED: "mdi:lock", +} + +NODE_STATUS_ICON: dict[NodeStatus, str] = { NodeStatus.ALIVE: "mdi:heart-pulse", NodeStatus.ASLEEP: "mdi:sleep", NodeStatus.AWAKE: "mdi:eye", @@ -529,6 +535,13 @@ def async_add_sensor(info: ZwaveDiscoveryInfo) -> None: async_add_entities(entities) + @callback + def async_add_controller_status_sensor() -> None: + """Add controller status sensor.""" + driver = client.driver + assert driver is not None # Driver is ready before platforms are loaded. + async_add_entities([ZWaveControllerStatusSensor(config_entry, driver)]) + @callback def async_add_node_status_sensor(node: ZwaveNode) -> None: """Add node status sensor.""" @@ -565,6 +578,14 @@ def async_add_statistics_sensors(node: ZwaveNode) -> None: ) ) + config_entry.async_on_unload( + async_dispatcher_connect( + hass, + f"{DOMAIN}_{config_entry.entry_id}_add_controller_status_sensor", + async_add_controller_status_sensor, + ) + ) + config_entry.async_on_unload( async_dispatcher_connect( hass, @@ -828,7 +849,7 @@ def _status_changed(self, _: dict) -> None: @property def icon(self) -> str | None: """Icon of the entity.""" - return STATUS_ICON[self.node.status] + return NODE_STATUS_ICON[self.node.status] async def async_added_to_hass(self) -> None: """Call when entity is added.""" @@ -856,6 +877,72 @@ async def async_added_to_hass(self) -> None: self.async_write_ha_state() +class ZWaveControllerStatusSensor(SensorEntity): + """Representation of a controller status sensor.""" + + _attr_should_poll = False + _attr_entity_category = EntityCategory.DIAGNOSTIC + _attr_has_entity_name = True + + def __init__(self, config_entry: ConfigEntry, driver: Driver) -> None: + """Initialize a generic Z-Wave device entity.""" + self.config_entry = config_entry + self.controller = driver.controller + node = self.controller.own_node + assert node + + # Entity class attributes + self._attr_name = "Status" + self._base_unique_id = get_valueless_base_unique_id(driver, node) + self._attr_unique_id = f"{self._base_unique_id}.controller_status" + # device may not be precreated in main handler yet + self._attr_device_info = get_device_info(driver, node) + + async def async_poll_value(self, _: bool) -> None: + """Poll a value.""" + # We log an error instead of raising an exception because this service call occurs + # in a separate task since it is called via the dispatcher and we don't want to + # raise the exception in that separate task because it is confusing to the user. + LOGGER.error( + "There is no value to refresh for this entity so the zwave_js.refresh_value" + " service won't work for it" + ) + + @callback + def _status_changed(self, _: dict) -> None: + """Call when status event is received.""" + self._attr_native_value = self.controller.status.name.lower() + self.async_write_ha_state() + + @property + def icon(self) -> str | None: + """Icon of the entity.""" + return CONTROLLER_STATUS_ICON[self.controller.status] + + async def async_added_to_hass(self) -> None: + """Call when entity is added.""" + # Add value_changed callbacks. + self.async_on_remove(self.controller.on("status changed", self._status_changed)) + self.async_on_remove( + async_dispatcher_connect( + self.hass, + f"{DOMAIN}_{self.unique_id}_poll_value", + self.async_poll_value, + ) + ) + # we don't listen for `remove_entity_on_ready_node` signal because this is not + # a regular node + self.async_on_remove( + async_dispatcher_connect( + self.hass, + f"{DOMAIN}_{self._base_unique_id}_remove_entity", + self.async_remove, + ) + ) + self._attr_native_value: str = self.controller.status.name.lower() + self.async_write_ha_state() + + class ZWaveStatisticsSensor(SensorEntity): """Representation of a node/controller statistics sensor.""" diff --git a/tests/components/zwave_js/fixtures/controller_state.json b/tests/components/zwave_js/fixtures/controller_state.json index 566ad3b6f2bac..d6d9dcacd9e25 100644 --- a/tests/components/zwave_js/fixtures/controller_state.json +++ b/tests/components/zwave_js/fixtures/controller_state.json @@ -24,7 +24,8 @@ "sucNodeId": 1, "supportsTimers": false, "isHealNetworkActive": false, - "inclusionState": 0 + "inclusionState": 0, + "status": 0 }, "nodes": [] } diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py index 5bafe9323625d..a717b58be2b55 100644 --- a/tests/components/zwave_js/test_api.py +++ b/tests/components/zwave_js/test_api.py @@ -244,6 +244,52 @@ async def test_network_status( assert msg["error"]["code"] == ERR_INVALID_FORMAT +async def test_subscribe_controller_status( + hass: HomeAssistant, + client, + integration, + hass_ws_client: WebSocketGenerator, +) -> None: + """Test the subscribe controller status websocket command.""" + entry = integration + ws_client = await hass_ws_client(hass) + driver = client.driver + node = driver.controller.nodes[1] + + dev_reg = dr.async_get(hass) + device = dev_reg.async_get_or_create( + config_entry_id=entry.entry_id, identifiers={get_device_id(driver, node)} + ) + + await ws_client.send_json( + { + ID: 3, + TYPE: "zwave_js/subscribe_controller_status", + DEVICE_ID: device.id, + } + ) + + msg = await ws_client.receive_json() + assert msg["success"] + assert msg["result"] == 0 + + event = Event( + "status changed", + { + "source": "controller", + "event": "status changed", + "status": 1, + }, + ) + driver.controller.receive_event(event) + await hass.async_block_till_done() + + msg = await ws_client.receive_json() + + assert msg["event"]["event"] == "status changed" + assert msg["event"]["status"] == 1 + + async def test_subscribe_node_status( hass: HomeAssistant, multisensor_6_state, diff --git a/tests/components/zwave_js/test_sensor.py b/tests/components/zwave_js/test_sensor.py index d809d52821c66..67c0d9769ecf8 100644 --- a/tests/components/zwave_js/test_sensor.py +++ b/tests/components/zwave_js/test_sensor.py @@ -261,6 +261,41 @@ async def test_config_parameter_sensor( await hass.async_block_till_done() +async def test_controller_status_sensor( + hass: HomeAssistant, client, integration +) -> None: + """Test controller status sensor is created and gets updated on controller state changes.""" + entity_id = "sensor.z_stick_gen5_usb_controller_status" + ent_reg = er.async_get(hass) + entity_entry = ent_reg.async_get(entity_id) + + assert not entity_entry.disabled + assert entity_entry.entity_category is EntityCategory.DIAGNOSTIC + assert hass.states.get(entity_id).state == "ready" + assert hass.states.get(entity_id).attributes[ATTR_ICON] == "mdi:check" + + event = Event( + "status changed", + data={"source": "controller", "event": "status changed", "status": 1}, + ) + client.driver.controller.receive_event(event) + assert hass.states.get(entity_id).state == "unresponsive" + assert hass.states.get(entity_id).attributes[ATTR_ICON] == "mdi:bell-off" + + # Test transitions work + event = Event( + "status changed", + data={"source": "controller", "event": "status changed", "status": 2}, + ) + client.driver.controller.receive_event(event) + assert hass.states.get(entity_id).state == "jammed" + assert hass.states.get(entity_id).attributes[ATTR_ICON] == "mdi:lock" + + # Disconnect the client and make sure the entity is still available + await client.disconnect() + assert hass.states.get(entity_id).state != STATE_UNAVAILABLE + + async def test_node_status_sensor( hass: HomeAssistant, client, lock_id_lock_as_id150, integration ) -> None: @@ -325,6 +360,16 @@ async def test_node_status_sensor( is None ) + # Assert a controller status sensor entity is not created for a node + assert ( + ent_reg.async_get_entity_id( + DOMAIN, + "sensor", + f"{get_valueless_base_unique_id(driver, node)}.controller_status", + ) + is None + ) + async def test_node_status_sensor_not_ready( hass: HomeAssistant, From e63845919df0532b4b5ae963a31544a74c6b6426 Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Tue, 29 Aug 2023 02:17:31 -0400 Subject: [PATCH 2/7] Also update network status command --- homeassistant/components/zwave_js/api.py | 1 + 1 file changed, 1 insertion(+) diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py index 7d1a9a40c3db7..3718eafaa4f48 100644 --- a/homeassistant/components/zwave_js/api.py +++ b/homeassistant/components/zwave_js/api.py @@ -515,6 +515,7 @@ async def websocket_network_status( "is_heal_network_active": controller.is_heal_network_active, "inclusion_state": controller.inclusion_state, "rf_region": controller.rf_region, + "status": controller.status, "nodes": [node_status(node) for node in driver.controller.nodes.values()], }, } From cd4149c5bc363d44af7b9b0aaf7cf37c4706c293 Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Tue, 29 Aug 2023 03:19:53 -0400 Subject: [PATCH 3/7] fix tests --- tests/components/zwave_js/test_discovery.py | 4 +++- tests/components/zwave_js/test_init.py | 10 +++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/components/zwave_js/test_discovery.py b/tests/components/zwave_js/test_discovery.py index 99a46eaadf913..cbaa27c2a9130 100644 --- a/tests/components/zwave_js/test_discovery.py +++ b/tests/components/zwave_js/test_discovery.py @@ -227,7 +227,9 @@ async def test_indicator_test( assert len(hass.states.async_entity_ids(NUMBER_DOMAIN)) == 0 assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 1 # only ping assert len(hass.states.async_entity_ids(BINARY_SENSOR_DOMAIN)) == 1 - assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2 # include node status + assert ( + len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3 + ) # include node + controller status assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 entity_id = "binary_sensor.this_is_a_fake_device_binary_sensor" diff --git a/tests/components/zwave_js/test_init.py b/tests/components/zwave_js/test_init.py index c421e0434130d..4d9972297fc34 100644 --- a/tests/components/zwave_js/test_init.py +++ b/tests/components/zwave_js/test_init.py @@ -204,7 +204,7 @@ async def test_on_node_added_not_ready( dev_reg = dr.async_get(hass) device_id = f"{client.driver.controller.home_id}-{zp3111_not_ready_state['nodeId']}" - assert len(hass.states.async_all()) == 0 + assert len(hass.states.async_all()) == 1 assert len(dev_reg.devices) == 1 node_state = deepcopy(zp3111_not_ready_state) @@ -223,7 +223,7 @@ async def test_on_node_added_not_ready( await hass.async_block_till_done() # the only entities are the node status sensor and ping button - assert len(hass.states.async_all()) == 2 + assert len(hass.states.async_all()) == 3 device = dev_reg.async_get_device(identifiers={(DOMAIN, device_id)}) assert device @@ -325,7 +325,7 @@ async def test_existing_node_not_ready( assert not device.sw_version # the only entities are the node status sensor and ping button - assert len(hass.states.async_all()) == 2 + assert len(hass.states.async_all()) == 3 device = dev_reg.async_get_device(identifiers={(DOMAIN, device_id)}) assert device @@ -963,7 +963,7 @@ async def test_removed_device( # Check how many entities there are ent_reg = er.async_get(hass) entity_entries = er.async_entries_for_config_entry(ent_reg, integration.entry_id) - assert len(entity_entries) == 91 + assert len(entity_entries) == 92 # Remove a node and reload the entry old_node = driver.controller.nodes.pop(13) @@ -975,7 +975,7 @@ async def test_removed_device( device_entries = dr.async_entries_for_config_entry(dev_reg, integration.entry_id) assert len(device_entries) == 2 entity_entries = er.async_entries_for_config_entry(ent_reg, integration.entry_id) - assert len(entity_entries) == 60 + assert len(entity_entries) == 61 assert ( dev_reg.async_get_device(identifiers={get_device_id(driver, old_node)}) is None ) From 9b0183f7ac52813aa89c106ce15f483f40cec6e8 Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Wed, 30 Aug 2023 02:51:53 -0400 Subject: [PATCH 4/7] Remove WS command since we have a sensor entity --- homeassistant/components/zwave_js/api.py | 42 ---------------------- tests/components/zwave_js/test_api.py | 46 ------------------------ 2 files changed, 88 deletions(-) diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py index 418852621caa9..d93745f7a6664 100644 --- a/homeassistant/components/zwave_js/api.py +++ b/homeassistant/components/zwave_js/api.py @@ -391,7 +391,6 @@ def node_status(node: Node) -> dict[str, Any]: def async_register_api(hass: HomeAssistant) -> None: """Register all of our api endpoints.""" websocket_api.async_register_command(hass, websocket_network_status) - websocket_api.async_register_command(hass, websocket_subscribe_controller_status) websocket_api.async_register_command(hass, websocket_subscribe_node_status) websocket_api.async_register_command(hass, websocket_node_status) websocket_api.async_register_command(hass, websocket_node_metadata) @@ -525,47 +524,6 @@ async def websocket_network_status( ) -@websocket_api.websocket_command( - { - vol.Required(TYPE): "zwave_js/subscribe_controller_status", - vol.Required(DEVICE_ID): str, - } -) -@websocket_api.async_response -@async_get_node -async def websocket_subscribe_controller_status( - hass: HomeAssistant, - connection: ActiveConnection, - msg: dict[str, Any], - node: Node, -) -> None: - """Subscribe to controller status update events of a Z-Wave JS controller.""" - driver = node.client.driver - assert driver - controller = driver.controller - - @callback - def forward_event(event: dict) -> None: - """Forward the event.""" - connection.send_message( - websocket_api.event_message( - msg[ID], - {"event": event["event"], "status": controller.status}, - ) - ) - - @callback - def async_cleanup() -> None: - """Remove signal listeners.""" - for unsub in unsubs: - unsub() - - connection.subscriptions[msg[ID]] = async_cleanup - msg[DATA_UNSUBSCRIBE] = unsubs = [controller.on("status changed", forward_event)] - - connection.send_result(msg[ID], controller.status) - - @websocket_api.websocket_command( { vol.Required(TYPE): "zwave_js/subscribe_node_status", diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py index 284a1e5d1f3f7..e686def8883cb 100644 --- a/tests/components/zwave_js/test_api.py +++ b/tests/components/zwave_js/test_api.py @@ -244,52 +244,6 @@ async def test_network_status( assert msg["error"]["code"] == ERR_INVALID_FORMAT -async def test_subscribe_controller_status( - hass: HomeAssistant, - client, - integration, - hass_ws_client: WebSocketGenerator, -) -> None: - """Test the subscribe controller status websocket command.""" - entry = integration - ws_client = await hass_ws_client(hass) - driver = client.driver - node = driver.controller.nodes[1] - - dev_reg = dr.async_get(hass) - device = dev_reg.async_get_or_create( - config_entry_id=entry.entry_id, identifiers={get_device_id(driver, node)} - ) - - await ws_client.send_json( - { - ID: 3, - TYPE: "zwave_js/subscribe_controller_status", - DEVICE_ID: device.id, - } - ) - - msg = await ws_client.receive_json() - assert msg["success"] - assert msg["result"] == 0 - - event = Event( - "status changed", - { - "source": "controller", - "event": "status changed", - "status": 1, - }, - ) - driver.controller.receive_event(event) - await hass.async_block_till_done() - - msg = await ws_client.receive_json() - - assert msg["event"]["event"] == "status changed" - assert msg["event"]["status"] == 1 - - async def test_subscribe_node_status( hass: HomeAssistant, multisensor_6_state, From 4cc87d99e18b935f27e7f9a4747f8a7ad1f21390 Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Wed, 30 Aug 2023 06:33:37 -0400 Subject: [PATCH 5/7] Update sensor.py Co-authored-by: Martin Hjelmare --- homeassistant/components/zwave_js/sensor.py | 1 - 1 file changed, 1 deletion(-) diff --git a/homeassistant/components/zwave_js/sensor.py b/homeassistant/components/zwave_js/sensor.py index cdc4deb5329d3..e3d5de1248e85 100644 --- a/homeassistant/components/zwave_js/sensor.py +++ b/homeassistant/components/zwave_js/sensor.py @@ -940,7 +940,6 @@ async def async_added_to_hass(self) -> None: ) ) self._attr_native_value: str = self.controller.status.name.lower() - self.async_write_ha_state() class ZWaveStatisticsSensor(SensorEntity): From d888899e34050d241f8f4c497c492a02a402ad36 Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Wed, 30 Aug 2023 06:39:33 -0400 Subject: [PATCH 6/7] move driver assertion out of closures --- homeassistant/components/zwave_js/sensor.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/zwave_js/sensor.py b/homeassistant/components/zwave_js/sensor.py index e3d5de1248e85..3c22288a1d696 100644 --- a/homeassistant/components/zwave_js/sensor.py +++ b/homeassistant/components/zwave_js/sensor.py @@ -491,12 +491,12 @@ async def async_setup_entry( ) -> None: """Set up Z-Wave sensor from config entry.""" client: ZwaveClient = hass.data[DOMAIN][config_entry.entry_id][DATA_CLIENT] + driver = client.driver + assert driver is not None # Driver is ready before platforms are loaded. @callback def async_add_sensor(info: ZwaveDiscoveryInfo) -> None: """Add Z-Wave Sensor.""" - driver = client.driver - assert driver is not None # Driver is ready before platforms are loaded. entities: list[ZWaveBaseEntity] = [] if info.platform_data: @@ -538,22 +538,16 @@ def async_add_sensor(info: ZwaveDiscoveryInfo) -> None: @callback def async_add_controller_status_sensor() -> None: """Add controller status sensor.""" - driver = client.driver - assert driver is not None # Driver is ready before platforms are loaded. async_add_entities([ZWaveControllerStatusSensor(config_entry, driver)]) @callback def async_add_node_status_sensor(node: ZwaveNode) -> None: """Add node status sensor.""" - driver = client.driver - assert driver is not None # Driver is ready before platforms are loaded. async_add_entities([ZWaveNodeStatusSensor(config_entry, driver, node)]) @callback def async_add_statistics_sensors(node: ZwaveNode) -> None: """Add statistics sensors.""" - driver = client.driver - assert driver is not None # Driver is ready before platforms are loaded. async_add_entities( [ ZWaveStatisticsSensor( From 7271a7485106d765808881952596814af96c6e60 Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Wed, 30 Aug 2023 06:41:26 -0400 Subject: [PATCH 7/7] store state in tests --- tests/components/zwave_js/test_sensor.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/components/zwave_js/test_sensor.py b/tests/components/zwave_js/test_sensor.py index 67c0d9769ecf8..d452f28b3bf5f 100644 --- a/tests/components/zwave_js/test_sensor.py +++ b/tests/components/zwave_js/test_sensor.py @@ -271,16 +271,20 @@ async def test_controller_status_sensor( assert not entity_entry.disabled assert entity_entry.entity_category is EntityCategory.DIAGNOSTIC - assert hass.states.get(entity_id).state == "ready" - assert hass.states.get(entity_id).attributes[ATTR_ICON] == "mdi:check" + state = hass.states.get(entity_id) + assert state + assert state.state == "ready" + assert state.attributes[ATTR_ICON] == "mdi:check" event = Event( "status changed", data={"source": "controller", "event": "status changed", "status": 1}, ) client.driver.controller.receive_event(event) - assert hass.states.get(entity_id).state == "unresponsive" - assert hass.states.get(entity_id).attributes[ATTR_ICON] == "mdi:bell-off" + state = hass.states.get(entity_id) + assert state + assert state.state == "unresponsive" + assert state.attributes[ATTR_ICON] == "mdi:bell-off" # Test transitions work event = Event( @@ -288,8 +292,10 @@ async def test_controller_status_sensor( data={"source": "controller", "event": "status changed", "status": 2}, ) client.driver.controller.receive_event(event) - assert hass.states.get(entity_id).state == "jammed" - assert hass.states.get(entity_id).attributes[ATTR_ICON] == "mdi:lock" + state = hass.states.get(entity_id) + assert state + assert state.state == "jammed" + assert state.attributes[ATTR_ICON] == "mdi:lock" # Disconnect the client and make sure the entity is still available await client.disconnect()