256 lines
8.7 KiB
Python
256 lines
8.7 KiB
Python
"""Common test objects."""
|
|
|
|
import asyncio
|
|
from datetime import timedelta
|
|
import math
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
import zigpy.zcl
|
|
import zigpy.zcl.foundation as zcl_f
|
|
|
|
import homeassistant.components.zha.core.const as zha_const
|
|
from homeassistant.components.zha.core.helpers import (
|
|
async_get_zha_config_value,
|
|
get_zha_gateway,
|
|
)
|
|
from homeassistant.core import HomeAssistant
|
|
from homeassistant.helpers import entity_registry as er
|
|
import homeassistant.util.dt as dt_util
|
|
|
|
from tests.common import async_fire_time_changed
|
|
|
|
|
|
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 = cluster.attributes.get(attr_id)
|
|
|
|
if attr is not None:
|
|
value = cluster.PLUGGED_ATTR_READS.get(attr.name)
|
|
if value is not None:
|
|
result.append(
|
|
zcl_f.ReadAttributeRecord(
|
|
attr_id,
|
|
zcl_f.Status.SUCCESS,
|
|
zcl_f.TypeValue(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.configure_reporting_multiple = AsyncMock(
|
|
return_value=zcl_f.ConfigureReportingResponse.deserialize(b"\x00")[0]
|
|
)
|
|
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(wraps=cluster.write_attributes)
|
|
cluster._write_attributes = AsyncMock(
|
|
return_value=[zcl_f.WriteAttributesResponse.deserialize(b"\x00")[0]]
|
|
)
|
|
if cluster.cluster_id == 4:
|
|
cluster.add = AsyncMock(return_value=[0])
|
|
|
|
|
|
def update_attribute_cache(cluster):
|
|
"""Update attribute cache based on plugged attributes."""
|
|
if not cluster.PLUGGED_ATTR_READS:
|
|
return
|
|
|
|
attrs = []
|
|
for attrid, value in cluster.PLUGGED_ATTR_READS.items():
|
|
if isinstance(attrid, str):
|
|
attrid = cluster.attributes_by_name[attrid].id
|
|
else:
|
|
attrid = zigpy.types.uint16_t(attrid)
|
|
attrs.append(make_attribute(attrid, value))
|
|
|
|
hdr = make_zcl_header(zcl_f.GeneralCommand.Report_Attributes)
|
|
hdr.frame_control.disable_default_response = True
|
|
msg = zcl_f.GENERAL_COMMANDS[zcl_f.GeneralCommand.Report_Attributes].schema(
|
|
attribute_reports=attrs
|
|
)
|
|
cluster.handle_message(hdr, msg)
|
|
|
|
|
|
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: HomeAssistant, cluster: zigpy.zcl.Cluster, 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 = []
|
|
|
|
for attrid, value in attributes.items():
|
|
if isinstance(attrid, str):
|
|
attrid = cluster.attributes_by_name[attrid].id
|
|
else:
|
|
attrid = zigpy.types.uint16_t(attrid)
|
|
|
|
attrs.append(make_attribute(attrid, value))
|
|
|
|
msg = zcl_f.GENERAL_COMMANDS[zcl_f.GeneralCommand.Report_Attributes].schema(
|
|
attribute_reports=attrs
|
|
)
|
|
|
|
hdr = make_zcl_header(zcl_f.GeneralCommand.Report_Attributes)
|
|
hdr.frame_control.disable_default_response = True
|
|
cluster.handle_message(hdr, msg)
|
|
await hass.async_block_till_done()
|
|
|
|
|
|
def find_entity_id(domain, zha_device, hass, qualifier=None):
|
|
"""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.
|
|
"""
|
|
entities = find_entity_ids(domain, zha_device, hass)
|
|
if not entities:
|
|
return None
|
|
if qualifier:
|
|
for entity_id in entities:
|
|
if qualifier in entity_id:
|
|
return entity_id
|
|
else:
|
|
return entities[0]
|
|
|
|
|
|
def find_entity_ids(domain, zha_device, hass):
|
|
"""Find the entity ids 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.
|
|
"""
|
|
|
|
registry = er.async_get(hass)
|
|
return [
|
|
entity.entity_id
|
|
for entity in er.async_entries_for_device(registry, zha_device.device_id)
|
|
if entity.domain == domain
|
|
]
|
|
|
|
|
|
def async_find_group_entity_id(hass, domain, group):
|
|
"""Find the group entity id under test."""
|
|
entity_id = f"{domain}.coordinator_manufacturer_coordinator_model_{group.name.lower().replace(' ', '_')}"
|
|
|
|
entity_ids = hass.states.async_entity_ids(domain)
|
|
assert entity_id in entity_ids
|
|
return entity_id
|
|
|
|
|
|
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.configure_reporting_multiple.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, strict=False):
|
|
assert cluster.bind.call_count == 1
|
|
assert cluster.bind.await_count == 1
|
|
if reports:
|
|
assert cluster.configure_reporting.call_count == 0
|
|
assert cluster.configure_reporting.await_count == 0
|
|
assert cluster.configure_reporting_multiple.call_count == math.ceil(
|
|
reports / zha_const.REPORT_CONFIG_ATTR_PER_REQ
|
|
)
|
|
assert cluster.configure_reporting_multiple.await_count == math.ceil(
|
|
reports / zha_const.REPORT_CONFIG_ATTR_PER_REQ
|
|
)
|
|
else:
|
|
# no reports at all
|
|
assert cluster.configure_reporting.call_count == reports
|
|
assert cluster.configure_reporting.await_count == reports
|
|
assert cluster.configure_reporting_multiple.call_count == reports
|
|
assert cluster.configure_reporting_multiple.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()
|
|
|
|
|
|
async def async_shift_time(hass):
|
|
"""Shift time to cause call later tasks to run."""
|
|
next_update = dt_util.utcnow() + timedelta(seconds=11)
|
|
async_fire_time_changed(hass, next_update)
|
|
await hass.async_block_till_done()
|
|
|
|
|
|
def patch_zha_config(component: str, overrides: dict[tuple[str, str], Any]):
|
|
"""Patch the ZHA custom configuration defaults."""
|
|
|
|
def new_get_config(config_entry, section, config_key, default):
|
|
if (section, config_key) in overrides:
|
|
return overrides[section, config_key]
|
|
return async_get_zha_config_value(config_entry, section, config_key, default)
|
|
|
|
return patch(
|
|
f"homeassistant.components.zha.{component}.async_get_zha_config_value",
|
|
side_effect=new_get_config,
|
|
)
|