253 lines
8.3 KiB
Python
253 lines
8.3 KiB
Python
"""Common test objects."""
|
|
import asyncio
|
|
import time
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import zigpy.device as zigpy_dev
|
|
import zigpy.endpoint as zigpy_ep
|
|
import zigpy.profiles.zha
|
|
import zigpy.types
|
|
import zigpy.zcl
|
|
import zigpy.zcl.clusters.general
|
|
import zigpy.zcl.foundation as zcl_f
|
|
import zigpy.zdo.types
|
|
|
|
import homeassistant.components.zha.core.const as zha_const
|
|
from homeassistant.util import slugify
|
|
|
|
|
|
class FakeEndpoint:
|
|
"""Fake endpoint for moking zigpy."""
|
|
|
|
def __init__(self, manufacturer, model, epid=1):
|
|
"""Init fake endpoint."""
|
|
self.device = None
|
|
self.endpoint_id = epid
|
|
self.in_clusters = {}
|
|
self.out_clusters = {}
|
|
self._cluster_attr = {}
|
|
self.member_of = {}
|
|
self.status = zigpy_ep.Status.ZDO_INIT
|
|
self.manufacturer = manufacturer
|
|
self.model = model
|
|
self.profile_id = zigpy.profiles.zha.PROFILE_ID
|
|
self.device_type = None
|
|
self.request = AsyncMock(return_value=[0])
|
|
|
|
def add_input_cluster(self, cluster_id, _patch_cluster=True):
|
|
"""Add an input cluster."""
|
|
cluster = zigpy.zcl.Cluster.from_id(self, cluster_id, is_server=True)
|
|
if _patch_cluster:
|
|
patch_cluster(cluster)
|
|
self.in_clusters[cluster_id] = cluster
|
|
ep_attribute = cluster.ep_attribute
|
|
if ep_attribute:
|
|
setattr(self, ep_attribute, cluster)
|
|
|
|
def add_output_cluster(self, cluster_id, _patch_cluster=True):
|
|
"""Add an output cluster."""
|
|
cluster = zigpy.zcl.Cluster.from_id(self, cluster_id, is_server=False)
|
|
if _patch_cluster:
|
|
patch_cluster(cluster)
|
|
self.out_clusters[cluster_id] = cluster
|
|
|
|
reply = AsyncMock(return_value=[0])
|
|
request = AsyncMock(return_value=[0])
|
|
|
|
@property
|
|
def __class__(self):
|
|
"""Fake being Zigpy endpoint."""
|
|
return zigpy_ep.Endpoint
|
|
|
|
@property
|
|
def unique_id(self):
|
|
"""Return the unique id for the endpoint."""
|
|
return self.device.ieee, self.endpoint_id
|
|
|
|
|
|
FakeEndpoint.add_to_group = zigpy_ep.Endpoint.add_to_group
|
|
FakeEndpoint.remove_from_group = zigpy_ep.Endpoint.remove_from_group
|
|
|
|
|
|
def patch_cluster(cluster):
|
|
"""Patch a cluster for testing."""
|
|
cluster.PLUGGED_ATTR_READS = {}
|
|
|
|
async def _read_attribute_raw(attributes, *args, **kwargs):
|
|
result = []
|
|
for attr_id in attributes:
|
|
value = cluster.PLUGGED_ATTR_READS.get(attr_id)
|
|
if value is None:
|
|
# try converting attr_id to attr_name and lookup the plugs again
|
|
attr_name = cluster.attributes.get(attr_id)
|
|
value = attr_name and cluster.PLUGGED_ATTR_READS.get(attr_name[0])
|
|
if value is not None:
|
|
result.append(
|
|
zcl_f.ReadAttributeRecord(
|
|
attr_id,
|
|
zcl_f.Status.SUCCESS,
|
|
zcl_f.TypeValue(python_type=None, value=value),
|
|
)
|
|
)
|
|
else:
|
|
result.append(zcl_f.ReadAttributeRecord(attr_id, zcl_f.Status.FAILURE))
|
|
return (result,)
|
|
|
|
cluster.bind = AsyncMock(return_value=[0])
|
|
cluster.configure_reporting = AsyncMock(
|
|
return_value=[
|
|
[zcl_f.ConfigureReportingResponseRecord(zcl_f.Status.SUCCESS, 0x00, 0xAABB)]
|
|
]
|
|
)
|
|
cluster.deserialize = Mock()
|
|
cluster.handle_cluster_request = Mock()
|
|
cluster.read_attributes = AsyncMock(wraps=cluster.read_attributes)
|
|
cluster.read_attributes_raw = AsyncMock(side_effect=_read_attribute_raw)
|
|
cluster.unbind = AsyncMock(return_value=[0])
|
|
cluster.write_attributes = AsyncMock(
|
|
return_value=[zcl_f.WriteAttributesResponse.deserialize(b"\x00")[0]]
|
|
)
|
|
if cluster.cluster_id == 4:
|
|
cluster.add = AsyncMock(return_value=[0])
|
|
|
|
|
|
class FakeDevice:
|
|
"""Fake device for mocking zigpy."""
|
|
|
|
def __init__(self, app, ieee, manufacturer, model, node_desc=None, nwk=0xB79C):
|
|
"""Init fake device."""
|
|
self._application = app
|
|
self.application = app
|
|
self.ieee = zigpy.types.EUI64.convert(ieee)
|
|
self.nwk = nwk
|
|
self.zdo = Mock()
|
|
self.endpoints = {0: self.zdo}
|
|
self.lqi = 255
|
|
self.rssi = 8
|
|
self.last_seen = time.time()
|
|
self.status = zigpy_dev.Status.ENDPOINTS_INIT
|
|
self.initializing = False
|
|
self.skip_configuration = False
|
|
self.manufacturer = manufacturer
|
|
self.model = model
|
|
self.remove_from_group = AsyncMock()
|
|
if node_desc is None:
|
|
node_desc = b"\x02@\x807\x10\x7fd\x00\x00*d\x00\x00"
|
|
self.node_desc = zigpy.zdo.types.NodeDescriptor.deserialize(node_desc)[0]
|
|
self.neighbors = []
|
|
|
|
|
|
FakeDevice.add_to_group = zigpy_dev.Device.add_to_group
|
|
|
|
|
|
def get_zha_gateway(hass):
|
|
"""Return ZHA gateway from hass.data."""
|
|
try:
|
|
return hass.data[zha_const.DATA_ZHA][zha_const.DATA_ZHA_GATEWAY]
|
|
except KeyError:
|
|
return None
|
|
|
|
|
|
def make_attribute(attrid, value, status=0):
|
|
"""Make an attribute."""
|
|
attr = zcl_f.Attribute()
|
|
attr.attrid = attrid
|
|
attr.value = zcl_f.TypeValue()
|
|
attr.value.value = value
|
|
return attr
|
|
|
|
|
|
def send_attribute_report(hass, cluster, attrid, value):
|
|
"""Send a single attribute report."""
|
|
return send_attributes_report(hass, cluster, {attrid: value})
|
|
|
|
|
|
async def send_attributes_report(hass, cluster: int, attributes: dict):
|
|
"""Cause the sensor to receive an attribute report from the network.
|
|
|
|
This is to simulate the normal device communication that happens when a
|
|
device is paired to the zigbee network.
|
|
"""
|
|
attrs = [make_attribute(attrid, value) for attrid, value in attributes.items()]
|
|
hdr = make_zcl_header(zcl_f.Command.Report_Attributes)
|
|
hdr.frame_control.disable_default_response = True
|
|
cluster.handle_message(hdr, [attrs])
|
|
await hass.async_block_till_done()
|
|
|
|
|
|
async def find_entity_id(domain, zha_device, hass):
|
|
"""Find the entity id under the testing.
|
|
|
|
This is used to get the entity id in order to get the state from the state
|
|
machine so that we can test state changes.
|
|
"""
|
|
ieeetail = "".join([f"{o:02x}" for o in zha_device.ieee[:4]])
|
|
head = f"{domain}.{slugify(f'{zha_device.name} {ieeetail}')}"
|
|
|
|
enitiy_ids = hass.states.async_entity_ids(domain)
|
|
await hass.async_block_till_done()
|
|
|
|
for entity_id in enitiy_ids:
|
|
if entity_id.startswith(head):
|
|
return entity_id
|
|
return None
|
|
|
|
|
|
def async_find_group_entity_id(hass, domain, group):
|
|
"""Find the group entity id under test."""
|
|
entity_id = f"{domain}.{group.name.lower().replace(' ','_')}_zha_group_0x{group.group_id:04x}"
|
|
|
|
entity_ids = hass.states.async_entity_ids(domain)
|
|
|
|
if entity_id in entity_ids:
|
|
return entity_id
|
|
return None
|
|
|
|
|
|
async def async_enable_traffic(hass, zha_devices, enabled=True):
|
|
"""Allow traffic to flow through the gateway and the zha device."""
|
|
for zha_device in zha_devices:
|
|
zha_device.update_available(enabled)
|
|
await hass.async_block_till_done()
|
|
|
|
|
|
def make_zcl_header(
|
|
command_id: int, global_command: bool = True, tsn: int = 1
|
|
) -> 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=tsn, command_id=command_id)
|
|
|
|
|
|
def reset_clusters(clusters):
|
|
"""Reset mocks on cluster."""
|
|
for cluster in clusters:
|
|
cluster.bind.reset_mock()
|
|
cluster.configure_reporting.reset_mock()
|
|
cluster.write_attributes.reset_mock()
|
|
|
|
|
|
async def async_test_rejoin(hass, zigpy_device, clusters, report_counts, ep_id=1):
|
|
"""Test device rejoins."""
|
|
reset_clusters(clusters)
|
|
|
|
zha_gateway = get_zha_gateway(hass)
|
|
await zha_gateway.async_device_initialized(zigpy_device)
|
|
await hass.async_block_till_done()
|
|
for cluster, reports in zip(clusters, report_counts):
|
|
assert cluster.bind.call_count == 1
|
|
assert cluster.bind.await_count == 1
|
|
assert cluster.configure_reporting.call_count == reports
|
|
assert cluster.configure_reporting.await_count == reports
|
|
|
|
|
|
async def async_wait_for_updates(hass):
|
|
"""Wait until all scheduled updates are executed."""
|
|
await hass.async_block_till_done()
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
await hass.async_block_till_done()
|