From 994794275af5aae7d16c52283a4b6af670acfe25 Mon Sep 17 00:00:00 2001 From: David Mulcahey Date: Fri, 4 Oct 2024 12:30:20 -0400 Subject: [PATCH] QuirkBuilder enhancements --- tests/test_quirks_v2.py | 78 +++++++++++++++++++++++++++++++++++++ zigpy/quirks/v2/__init__.py | 35 ++++++++++++++--- 2 files changed, 108 insertions(+), 5 deletions(-) diff --git a/tests/test_quirks_v2.py b/tests/test_quirks_v2.py index 7674b9847..1dcb0d0d5 100644 --- a/tests/test_quirks_v2.py +++ b/tests/test_quirks_v2.py @@ -158,6 +158,84 @@ class AttributeDefs(BaseAttributeDefs): # pylint: disable=too-few-public-method assert quirked not in registry +async def test_quirks_v2_model_manufacturer(device_mock): + """Test the potential exceptions when model and manufacturer are set up incorrectly.""" + registry = DeviceRegistry() + + with pytest.raises( + ValueError, + match="manufacturer and model must be provided together or completely omitted.", + ): + ( + QuirkBuilder(device_mock.manufacturer, model=None, registry=registry) + .adds(Basic.cluster_id) + .adds(OnOff.cluster_id) + .enum( + OnOff.AttributeDefs.start_up_on_off.name, + OnOff.StartUpOnOff, + OnOff.cluster_id, + ) + .add_to_registry() + ) + + with pytest.raises( + ValueError, + match="manufacturer and model must be provided together or completely omitted.", + ): + ( + QuirkBuilder(manufacturer=None, model=device_mock.model, registry=registry) + .adds(Basic.cluster_id) + .adds(OnOff.cluster_id) + .enum( + OnOff.AttributeDefs.start_up_on_off.name, + OnOff.StartUpOnOff, + OnOff.cluster_id, + ) + .add_to_registry() + ) + + with pytest.raises( + ValueError, + match="At least one manufacturer and model must be specified for a v2 quirk.", + ): + ( + QuirkBuilder(registry=registry) + .adds(Basic.cluster_id) + .adds(OnOff.cluster_id) + .enum( + OnOff.AttributeDefs.start_up_on_off.name, + OnOff.StartUpOnOff, + OnOff.cluster_id, + ) + .add_to_registry() + ) + + +async def test_quirks_v2_quirk_builder_cloning(device_mock): + """Test the quirk builder clone functionality.""" + registry = DeviceRegistry() + + base = ( + QuirkBuilder(registry=registry) + .adds(Basic.cluster_id) + .adds(OnOff.cluster_id) + .enum( + OnOff.AttributeDefs.start_up_on_off.name, + OnOff.StartUpOnOff, + OnOff.cluster_id, + ) + .applies_to("foo", "bar") + ) + + cloned = base.clone() + base.add_to_registry() + + cloned.applies_to(device_mock.manufacturer, device_mock.model).add_to_registry() + + quirked = registry.get_device(device_mock) + assert isinstance(quirked, CustomDeviceV2) + + async def test_quirks_v2_signature_match(device_mock): """Test the signature_matches filter.""" registry = DeviceRegistry() diff --git a/zigpy/quirks/v2/__init__.py b/zigpy/quirks/v2/__init__.py index 6362f145f..644b3dc8b 100644 --- a/zigpy/quirks/v2/__init__.py +++ b/zigpy/quirks/v2/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations import collections +from copy import deepcopy from enum import Enum import inspect import logging @@ -378,9 +379,17 @@ class QuirkBuilder: """Quirks V2 registry entry.""" def __init__( - self, manufacturer: str, model: str, registry: DeviceRegistry = _DEVICE_REGISTRY + self, + manufacturer: str | None = None, + model: str | None = None, + registry: DeviceRegistry = _DEVICE_REGISTRY, ) -> None: """Initialize the quirk builder.""" + if manufacturer and not model or model and not manufacturer: + raise ValueError( + "manufacturer and model must be provided together or completely omitted." + ) + self.registry: DeviceRegistry = registry self.manufacturer_model_metadata: list[ManufacturerModelMetadata] = [] self.filters: list[FilterType] = [] @@ -410,11 +419,13 @@ def __init__( self.quirk_file = pathlib.Path(caller.filename) self.quirk_file_line = caller.lineno - self.also_applies_to(manufacturer, model) + if manufacturer and model: + self.applies_to(manufacturer, model) + UNBUILT_QUIRK_BUILDERS.append(self) - def also_applies_to(self, manufacturer: str, model: str) -> QuirkBuilder: - """Register this quirks v2 entry for an additional manufacturer and model.""" + def applies_to(self, manufacturer: str, model: str) -> QuirkBuilder: + """Register this quirks v2 entry for the specified manufacturer and model.""" self.manufacturer_model_metadata.append( ManufacturerModelMetadata( # type: ignore[call-arg] manufacturer=manufacturer, model=model @@ -422,6 +433,9 @@ def also_applies_to(self, manufacturer: str, model: str) -> QuirkBuilder: ) return self + # backward compatibility + also_applies_to = applies_to + def filter(self, filter_function: FilterType) -> QuirkBuilder: """Add a filter and returns self. @@ -849,6 +863,10 @@ def device_automation_triggers( def add_to_registry(self) -> QuirksV2RegistryEntry: """Build the quirks v2 registry entry.""" + if not self.manufacturer_model_metadata: + raise ValueError( + "At least one manufacturer and model must be specified for a v2 quirk." + ) quirk: QuirksV2RegistryEntry = QuirksV2RegistryEntry( # type: ignore[call-arg] manufacturer_model_metadata=tuple(self.manufacturer_model_metadata), quirk_file=self.quirk_file, @@ -876,6 +894,13 @@ def add_to_registry(self) -> QuirksV2RegistryEntry: return quirk + def clone(self) -> QuirkBuilder: + """Clone this QuirkBuilder omitting manufacturer and model data.""" + new_builder = deepcopy(self) + new_builder.registry = self.registry + new_builder.manufacturer_model_metadata = [] + return new_builder + def add_to_registry_v2( manufacturer: str, model: str, registry: DeviceRegistry = _DEVICE_REGISTRY @@ -885,4 +910,4 @@ def add_to_registry_v2( "add_to_registry_v2 is deprecated and will be removed in a future release. " "Please QuirkBuilder() instead and ensure you call add_to_registry()." ) - return QuirkBuilder(manufacturer, model, registry=registry) + return QuirkBuilder(registry=registry).applies_to(manufacturer, model)