Leverage zigpy for IEEE address conversions (#27972)

* Refactor EUI64 conversions.

* Update ZHA dependencies.

* Update tests.

homeassistant/components/zha/api.py

@@ -4,6 +4,7 @@ import asyncio
 import logging
 
 import voluptuous as vol
+from zigpy.types.named import EUI64
 
 from homeassistant.components import websocket_api
 from homeassistant.core import callback
@@ -44,7 +45,7 @@ from .core.const import (
     WARNING_DEVICE_STROBE_HIGH,
     WARNING_DEVICE_STROBE_YES,
 )
-from .core.helpers import async_is_bindable_target, convert_ieee, get_matched_clusters
+from .core.helpers import async_is_bindable_target, get_matched_clusters
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -76,16 +77,16 @@ IEEE_SERVICE = "ieee_based_service"
 SERVICE_SCHEMAS = {
     SERVICE_PERMIT: vol.Schema(
         {
-            vol.Optional(ATTR_IEEE_ADDRESS, default=None): convert_ieee,
+            vol.Optional(ATTR_IEEE_ADDRESS, default=None): EUI64.convert,
             vol.Optional(ATTR_DURATION, default=60): vol.All(
                 vol.Coerce(int), vol.Range(0, 254)
             ),
         }
     ),
-    IEEE_SERVICE: vol.Schema({vol.Required(ATTR_IEEE_ADDRESS): convert_ieee}),
+    IEEE_SERVICE: vol.Schema({vol.Required(ATTR_IEEE_ADDRESS): EUI64.convert}),
     SERVICE_SET_ZIGBEE_CLUSTER_ATTRIBUTE: vol.Schema(
         {
-            vol.Required(ATTR_IEEE): convert_ieee,
+            vol.Required(ATTR_IEEE): EUI64.convert,
             vol.Required(ATTR_ENDPOINT_ID): cv.positive_int,
             vol.Required(ATTR_CLUSTER_ID): cv.positive_int,
             vol.Optional(ATTR_CLUSTER_TYPE, default=CLUSTER_TYPE_IN): cv.string,
@@ -96,7 +97,7 @@ SERVICE_SCHEMAS = {
     ),
     SERVICE_WARNING_DEVICE_SQUAWK: vol.Schema(
         {
-            vol.Required(ATTR_IEEE): convert_ieee,
+            vol.Required(ATTR_IEEE): EUI64.convert,
             vol.Optional(
                 ATTR_WARNING_DEVICE_MODE, default=WARNING_DEVICE_SQUAWK_MODE_ARMED
             ): cv.positive_int,
@@ -110,7 +111,7 @@ SERVICE_SCHEMAS = {
     ),
     SERVICE_WARNING_DEVICE_WARN: vol.Schema(
         {
-            vol.Required(ATTR_IEEE): convert_ieee,
+            vol.Required(ATTR_IEEE): EUI64.convert,
             vol.Optional(
                 ATTR_WARNING_DEVICE_MODE, default=WARNING_DEVICE_MODE_EMERGENCY
             ): cv.positive_int,
@@ -131,7 +132,7 @@ SERVICE_SCHEMAS = {
     ),
     SERVICE_ISSUE_ZIGBEE_CLUSTER_COMMAND: vol.Schema(
         {
-            vol.Required(ATTR_IEEE): convert_ieee,
+            vol.Required(ATTR_IEEE): EUI64.convert,
             vol.Required(ATTR_ENDPOINT_ID): cv.positive_int,
             vol.Required(ATTR_CLUSTER_ID): cv.positive_int,
             vol.Optional(ATTR_CLUSTER_TYPE, default=CLUSTER_TYPE_IN): cv.string,
@@ -149,7 +150,7 @@ SERVICE_SCHEMAS = {
 @websocket_api.websocket_command(
     {
         vol.Required("type"): "zha/devices/permit",
-        vol.Optional(ATTR_IEEE, default=None): convert_ieee,
+        vol.Optional(ATTR_IEEE, default=None): EUI64.convert,
         vol.Optional(ATTR_DURATION, default=60): vol.All(
             vol.Coerce(int), vol.Range(0, 254)
         ),
@@ -200,7 +201,7 @@ async def websocket_get_devices(hass, connection, msg):
 @websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command(
-    {vol.Required(TYPE): "zha/device", vol.Required(ATTR_IEEE): convert_ieee}
+    {vol.Required(TYPE): "zha/device", vol.Required(ATTR_IEEE): EUI64.convert}
 )
 async def websocket_get_device(hass, connection, msg):
     """Get ZHA devices."""
@@ -252,7 +253,7 @@ def async_get_device_info(hass, device, ha_device_registry=None):
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zha/devices/reconfigure",
-        vol.Required(ATTR_IEEE): convert_ieee,
+        vol.Required(ATTR_IEEE): EUI64.convert,
     }
 )
 async def websocket_reconfigure_node(hass, connection, msg):
@@ -267,7 +268,7 @@ async def websocket_reconfigure_node(hass, connection, msg):
 @websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command(
-    {vol.Required(TYPE): "zha/devices/clusters", vol.Required(ATTR_IEEE): convert_ieee}
+    {vol.Required(TYPE): "zha/devices/clusters", vol.Required(ATTR_IEEE): EUI64.convert}
 )
 async def websocket_device_clusters(hass, connection, msg):
     """Return a list of device clusters."""
@@ -305,7 +306,7 @@ async def websocket_device_clusters(hass, connection, msg):
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zha/devices/clusters/attributes",
-        vol.Required(ATTR_IEEE): convert_ieee,
+        vol.Required(ATTR_IEEE): EUI64.convert,
         vol.Required(ATTR_ENDPOINT_ID): int,
         vol.Required(ATTR_CLUSTER_ID): int,
         vol.Required(ATTR_CLUSTER_TYPE): str,
@@ -346,7 +347,7 @@ async def websocket_device_cluster_attributes(hass, connection, msg):
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zha/devices/clusters/commands",
-        vol.Required(ATTR_IEEE): convert_ieee,
+        vol.Required(ATTR_IEEE): EUI64.convert,
         vol.Required(ATTR_ENDPOINT_ID): int,
         vol.Required(ATTR_CLUSTER_ID): int,
         vol.Required(ATTR_CLUSTER_TYPE): str,
@@ -400,7 +401,7 @@ async def websocket_device_cluster_commands(hass, connection, msg):
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zha/devices/clusters/attributes/value",
-        vol.Required(ATTR_IEEE): convert_ieee,
+        vol.Required(ATTR_IEEE): EUI64.convert,
         vol.Required(ATTR_ENDPOINT_ID): int,
         vol.Required(ATTR_CLUSTER_ID): int,
         vol.Required(ATTR_CLUSTER_TYPE): str,
@@ -444,7 +445,7 @@ async def websocket_read_zigbee_cluster_attributes(hass, connection, msg):
 @websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command(
-    {vol.Required(TYPE): "zha/devices/bindable", vol.Required(ATTR_IEEE): convert_ieee}
+    {vol.Required(TYPE): "zha/devices/bindable", vol.Required(ATTR_IEEE): EUI64.convert}
 )
 async def websocket_get_bindable_devices(hass, connection, msg):
     """Directly bind devices."""
@@ -472,8 +473,8 @@ async def websocket_get_bindable_devices(hass, connection, msg):
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zha/devices/bind",
-        vol.Required(ATTR_SOURCE_IEEE): convert_ieee,
+        vol.Required(ATTR_SOURCE_IEEE): EUI64.convert,
-        vol.Required(ATTR_TARGET_IEEE): convert_ieee,
+        vol.Required(ATTR_TARGET_IEEE): EUI64.convert,
     }
 )
 async def websocket_bind_devices(hass, connection, msg):
@@ -494,8 +495,8 @@ async def websocket_bind_devices(hass, connection, msg):
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zha/devices/unbind",
-        vol.Required(ATTR_SOURCE_IEEE): convert_ieee,
+        vol.Required(ATTR_SOURCE_IEEE): EUI64.convert,
-        vol.Required(ATTR_TARGET_IEEE): convert_ieee,
+        vol.Required(ATTR_TARGET_IEEE): EUI64.convert,
     }
 )
 async def websocket_unbind_devices(hass, connection, msg):

homeassistant/components/zha/core/helpers.py

@@ -8,6 +8,8 @@ import asyncio
 import collections
 import logging
 
+from zigpy.types.named import EUI64
+
 from homeassistant.core import callback
 
 from .const import (
@@ -78,15 +80,6 @@ async def check_zigpy_connection(usb_path, radio_type, database_path):
     return True
 
 
-def convert_ieee(ieee_str):
-    """Convert given ieee string to EUI64."""
-    from zigpy.types import EUI64, uint8_t
-
-    if ieee_str is None:
-        return None
-    return EUI64([uint8_t(p, base=16) for p in ieee_str.split(":")])
-
-
 def get_attr_id_by_name(cluster, attr_name):
     """Get the attribute id for a cluster attribute by its name."""
     return next(
@@ -145,7 +138,7 @@ async def async_get_zha_device(hass, device_id):
     registry_device = device_registry.async_get(device_id)
     zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
     ieee_address = list(list(registry_device.identifiers)[0])[1]
-    ieee = convert_ieee(ieee_address)
+    ieee = EUI64.convert(ieee_address)
     return zha_gateway.devices[ieee]
 
 

homeassistant/components/zha/manifest.json

@@ -6,10 +6,10 @@
   "requirements": [
     "bellows-homeassistant==0.10.0",
     "zha-quirks==0.0.26",
-    "zigpy-deconz==0.5.0",
+    "zigpy-deconz==0.6.0",
-    "zigpy-homeassistant==0.9.0",
+    "zigpy-homeassistant==0.10.0",
-    "zigpy-xbee-homeassistant==0.5.0",
+    "zigpy-xbee-homeassistant==0.6.0",
-    "zigpy-zigate==0.4.1"
+    "zigpy-zigate==0.5.0"
   ],
   "dependencies": [],
   "codeowners": ["@dmulcahey", "@adminiuga"]

requirements_all.txt

@@ -2032,16 +2032,16 @@ zhong_hong_hvac==1.0.9
 ziggo-mediabox-xl==1.1.0
 
 # homeassistant.components.zha
-zigpy-deconz==0.5.0
+zigpy-deconz==0.6.0
 
 # homeassistant.components.zha
-zigpy-homeassistant==0.9.0
+zigpy-homeassistant==0.10.0
 
 # homeassistant.components.zha
-zigpy-xbee-homeassistant==0.5.0
+zigpy-xbee-homeassistant==0.6.0
 
 # homeassistant.components.zha
-zigpy-zigate==0.4.1
+zigpy-zigate==0.5.0
 
 # homeassistant.components.zoneminder
 zm-py==0.3.3

requirements_test_all.txt

@@ -641,13 +641,13 @@ zeroconf==0.23.0
 zha-quirks==0.0.26
 
 # homeassistant.components.zha
-zigpy-deconz==0.5.0
+zigpy-deconz==0.6.0
 
 # homeassistant.components.zha
-zigpy-homeassistant==0.9.0
+zigpy-homeassistant==0.10.0
 
 # homeassistant.components.zha
-zigpy-xbee-homeassistant==0.5.0
+zigpy-xbee-homeassistant==0.6.0
 
 # homeassistant.components.zha
-zigpy-zigate==0.4.1
+zigpy-zigate==0.5.0

tests/components/zha/common.py

@@ -3,6 +3,8 @@ import time
 from unittest.mock import Mock, patch
 
 from asynctest import CoroutineMock
+from zigpy.types.named import EUI64
+import zigpy.zcl.foundation as zcl_f
 
 from homeassistant.components.zha.core.const import (
     DATA_ZHA,
@@ -10,7 +12,6 @@ from homeassistant.components.zha.core.const import (
     DATA_ZHA_CONFIG,
     DATA_ZHA_DISPATCHERS,
 )
-from homeassistant.components.zha.core.helpers import convert_ieee
 from homeassistant.util import slugify
 
 from tests.common import mock_coro
@@ -21,7 +22,7 @@ class FakeApplication:
 
     def __init__(self):
         """Init fake application."""
-        self.ieee = convert_ieee("00:15:8d:00:02:32:4f:32")
+        self.ieee = EUI64.convert("00:15:8d:00:02:32:4f:32")
         self.nwk = 0x087D
 
 
@@ -71,7 +72,6 @@ def patch_cluster(cluster):
     cluster.configure_reporting = CoroutineMock(return_value=[0])
     cluster.deserialize = Mock()
     cluster.handle_cluster_request = Mock()
-    cluster.handle_cluster_general_request = Mock()
     cluster.read_attributes = CoroutineMock()
     cluster.read_attributes_raw = Mock()
     cluster.unbind = CoroutineMock(return_value=[0])
@@ -83,7 +83,7 @@ class FakeDevice:
     def __init__(self, ieee, manufacturer, model):
         """Init fake device."""
         self._application = APPLICATION
-        self.ieee = convert_ieee(ieee)
+        self.ieee = EUI64.convert(ieee)
         self.nwk = 0xB79C
         self.zdo = Mock()
         self.endpoints = {0: self.zdo}
@@ -230,3 +230,12 @@ async def async_test_device_join(
                 domain, zigpy_device, cluster, use_suffix=device_type is None
             )
             assert hass.states.get(entity_id) is not None
+
+
+def make_zcl_header(command_id: int, global_command: bool = True) -> zcl_f.ZCLHeader:
+    """Cluster.handle_message() ZCL Header helper."""
+    if global_command:
+        frc = zcl_f.FrameControl(zcl_f.FrameType.GLOBAL_COMMAND)
+    else:
+        frc = zcl_f.FrameControl(zcl_f.FrameType.CLUSTER_COMMAND)
+    return zcl_f.ZCLHeader(frc, tsn=1, command_id=command_id)

tests/components/zha/test_binary_sensor.py

@@ -1,12 +1,16 @@
 """Test zha binary sensor."""
+from zigpy.zcl.foundation import Command
+
 from homeassistant.components.binary_sensor import DOMAIN
-from homeassistant.const import STATE_ON, STATE_OFF, STATE_UNAVAILABLE
+from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE
+
 from .common import (
+    async_enable_traffic,
     async_init_zigpy_device,
+    async_test_device_join,
     make_attribute,
     make_entity_id,
-    async_test_device_join,
+    make_zcl_header,
-    async_enable_traffic,
 )
 
 
@@ -74,13 +78,15 @@ async def async_test_binary_sensor_on_off(hass, cluster, entity_id):
     """Test getting on and off messages for binary sensors."""
     # binary sensor on
     attr = make_attribute(0, 1)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_ON
 
     # binary sensor off
     attr.value.value = 0
-    cluster.handle_message(0, 0x0A, [[attr]])
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_OFF
 

tests/components/zha/test_device_tracker.py

@@ -1,19 +1,25 @@
 """Test ZHA Device Tracker."""
 from datetime import timedelta
 import time
+
+from zigpy.zcl.foundation import Command
+
 from homeassistant.components.device_tracker import DOMAIN, SOURCE_TYPE_ROUTER
-from homeassistant.const import STATE_HOME, STATE_NOT_HOME, STATE_UNAVAILABLE
 from homeassistant.components.zha.core.registries import (
     SMARTTHINGS_ARRIVAL_SENSOR_DEVICE_TYPE,
 )
+from homeassistant.const import STATE_HOME, STATE_NOT_HOME, STATE_UNAVAILABLE
 import homeassistant.util.dt as dt_util
+
 from .common import (
+    async_enable_traffic,
     async_init_zigpy_device,
+    async_test_device_join,
     make_attribute,
     make_entity_id,
-    async_test_device_join,
+    make_zcl_header,
-    async_enable_traffic,
 )
+
 from tests.common import async_fire_time_changed
 
 
@@ -67,10 +73,11 @@ async def test_device_tracker(hass, config_entry, zha_gateway):
 
     # turn state flip
     attr = make_attribute(0x0020, 23)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+    cluster.handle_message(hdr, [[attr]])
 
     attr = make_attribute(0x0021, 200)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    cluster.handle_message(hdr, [[attr]])
 
     zigpy_device.last_seen = time.time() + 10
     next_update = dt_util.utcnow() + timedelta(seconds=30)

tests/components/zha/test_fan.py

@@ -1,18 +1,30 @@
 """Test zha fan."""
 from unittest.mock import call, patch
+
+from zigpy.zcl.foundation import Command
+
 from homeassistant.components import fan
-from homeassistant.const import STATE_ON, STATE_OFF, STATE_UNAVAILABLE
 from homeassistant.components.fan import ATTR_SPEED, DOMAIN, SERVICE_SET_SPEED
-from homeassistant.const import ATTR_ENTITY_ID, SERVICE_TURN_ON, SERVICE_TURN_OFF
+from homeassistant.const import (
-from tests.common import mock_coro
+    ATTR_ENTITY_ID,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    STATE_OFF,
+    STATE_ON,
+    STATE_UNAVAILABLE,
+)
+
 from .common import (
+    async_enable_traffic,
     async_init_zigpy_device,
+    async_test_device_join,
     make_attribute,
     make_entity_id,
-    async_test_device_join,
+    make_zcl_header,
-    async_enable_traffic,
 )
 
+from tests.common import mock_coro
+
 
 async def test_fan(hass, config_entry, zha_gateway):
     """Test zha fan platform."""
@@ -44,13 +56,14 @@ async def test_fan(hass, config_entry, zha_gateway):
 
     # turn on at fan
     attr = make_attribute(0, 1)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_ON
 
     # turn off at fan
     attr.value.value = 0
-    cluster.handle_message(0, 0x0A, [[attr]])
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_OFF
 

tests/components/zha/test_light.py

@@ -2,6 +2,8 @@
 import asyncio
 from unittest.mock import MagicMock, call, patch, sentinel
 
+from zigpy.zcl.foundation import Command
+
 from homeassistant.components.light import DOMAIN
 from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE
 
@@ -11,6 +13,7 @@ from .common import (
     async_test_device_join,
     make_attribute,
     make_entity_id,
+    make_zcl_header,
 )
 
 from tests.common import mock_coro
@@ -123,13 +126,14 @@ async def async_test_on_off_from_light(hass, cluster, entity_id):
     """Test on off functionality from the light."""
     # turn on at light
     attr = make_attribute(0, 1)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_ON
 
     # turn off at light
     attr.value.value = 0
-    cluster.handle_message(0, 0x0A, [[attr]])
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_OFF
 
@@ -138,7 +142,8 @@ async def async_test_on_from_light(hass, cluster, entity_id):
     """Test on off functionality from the light."""
     # turn on at light
     attr = make_attribute(0, 1)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_ON
 
@@ -243,7 +248,8 @@ async def async_test_level_on_off_from_hass(
 async def async_test_dimmer_from_light(hass, cluster, entity_id, level, expected_state):
     """Test dimmer functionality from the light."""
     attr = make_attribute(0, level)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == expected_state
     # hass uses None for brightness of 0 in state attributes

tests/components/zha/test_lock.py

@@ -1,15 +1,21 @@
 """Test zha lock."""
 from unittest.mock import patch
-from homeassistant.const import STATE_LOCKED, STATE_UNLOCKED, STATE_UNAVAILABLE
+
+from zigpy.zcl.foundation import Command
+
 from homeassistant.components.lock import DOMAIN
-from tests.common import mock_coro
+from homeassistant.const import STATE_LOCKED, STATE_UNAVAILABLE, STATE_UNLOCKED
+
 from .common import (
+    async_enable_traffic,
     async_init_zigpy_device,
     make_attribute,
     make_entity_id,
-    async_enable_traffic,
+    make_zcl_header,
 )
 
+from tests.common import mock_coro
+
 LOCK_DOOR = 0
 UNLOCK_DOOR = 1
 
@@ -43,13 +49,14 @@ async def test_lock(hass, config_entry, zha_gateway):
 
     # set state to locked
     attr = make_attribute(0, 1)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_LOCKED
 
     # set state to unlocked
     attr.value.value = 2
-    cluster.handle_message(0, 0x0A, [[attr]])
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_UNLOCKED
 

tests/components/zha/test_sensor.py

@@ -1,12 +1,16 @@
 """Test zha sensor."""
+from zigpy.zcl.foundation import Command
+
 from homeassistant.components.sensor import DOMAIN
-from homeassistant.const import STATE_UNKNOWN, STATE_UNAVAILABLE
+from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN
+
 from .common import (
+    async_enable_traffic,
     async_init_zigpy_device,
+    async_test_device_join,
     make_attribute,
     make_entity_id,
-    async_test_device_join,
+    make_zcl_header,
-    async_enable_traffic,
 )
 
 
@@ -177,7 +181,8 @@ async def send_attribute_report(hass, cluster, attrid, value):
     device is paired to the zigbee network.
     """
     attr = make_attribute(attrid, value)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
 
 

tests/components/zha/test_switch.py

@@ -1,16 +1,22 @@
 """Test zha switch."""
 from unittest.mock import call, patch
+
+from zigpy.zcl.foundation import Command
+
 from homeassistant.components.switch import DOMAIN
-from homeassistant.const import STATE_ON, STATE_OFF, STATE_UNAVAILABLE
+from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE
-from tests.common import mock_coro
+
 from .common import (
+    async_enable_traffic,
     async_init_zigpy_device,
+    async_test_device_join,
     make_attribute,
     make_entity_id,
-    async_test_device_join,
+    make_zcl_header,
-    async_enable_traffic,
 )
 
+from tests.common import mock_coro
+
 ON = 1
 OFF = 0
 
@@ -44,13 +50,14 @@ async def test_switch(hass, config_entry, zha_gateway):
 
     # turn on at switch
     attr = make_attribute(0, 1)
-    cluster.handle_message(1, 0x0A, [[attr]])
+    hdr = make_zcl_header(Command.Report_Attributes)
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_ON
 
     # turn off at switch
     attr.value.value = 0
-    cluster.handle_message(0, 0x0A, [[attr]])
+    cluster.handle_message(hdr, [[attr]])
     await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_OFF