diff --git a/script/scaffold/templates/device_action/tests/test_device_action.py b/script/scaffold/templates/device_action/tests/test_device_action.py index d4116cdae84..16b54fb8a22 100644 --- a/script/scaffold/templates/device_action/tests/test_device_action.py +++ b/script/scaffold/templates/device_action/tests/test_device_action.py @@ -5,7 +5,7 @@ from homeassistant.components import automation from homeassistant.components.NEW_DOMAIN import DOMAIN from homeassistant.components.device_automation import DeviceAutomationType from homeassistant.core import HomeAssistant -from homeassistant.helpers import device_registry, entity_registry +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.entity import EntityCategory from homeassistant.setup import async_setup_component @@ -14,36 +14,24 @@ from tests.common import ( assert_lists_same, async_get_device_automations, async_mock_service, - mock_device_registry, - mock_registry, ) -@pytest.fixture -def device_reg(hass: HomeAssistant) -> device_registry.DeviceRegistry: - """Return an empty, loaded, registry.""" - return mock_device_registry(hass) - - -@pytest.fixture -def entity_reg(hass: HomeAssistant) -> entity_registry.EntityRegistry: - """Return an empty, loaded, registry.""" - return mock_registry(hass) - - async def test_get_actions( hass: HomeAssistant, - device_reg: device_registry.DeviceRegistry, - entity_reg: entity_registry.EntityRegistry, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, ) -> None: """Test we get the expected actions from a NEW_DOMAIN.""" config_entry = MockConfigEntry(domain="test", data={}) config_entry.add_to_hass(hass) - device_entry = device_reg.async_get_or_create( + device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + entity_registry.async_get_or_create( + DOMAIN, "test", "5678", device_id=device_entry.id ) - entity_reg.async_get_or_create(DOMAIN, "test", "5678", device_id=device_entry.id) expected_actions = [ { "domain": DOMAIN, @@ -62,27 +50,27 @@ async def test_get_actions( @pytest.mark.parametrize( "hidden_by,entity_category", ( - (entity_registry.RegistryEntryHider.INTEGRATION, None), - (entity_registry.RegistryEntryHider.USER, None), + (er.RegistryEntryHider.INTEGRATION, None), + (er.RegistryEntryHider.USER, None), (None, EntityCategory.CONFIG), (None, EntityCategory.DIAGNOSTIC), ), ) async def test_get_actions_hidden_auxiliary( - hass, - device_reg, - entity_reg, - hidden_by, - entity_category, + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + hidden_by: er.RegistryEntryHider | None, + entity_category: EntityCategory | None, ): """Test we get the expected actions from a hidden or auxiliary entity.""" config_entry = MockConfigEntry(domain="test", data={}) config_entry.add_to_hass(hass) - device_entry = device_reg.async_get_or_create( + device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, ) - entity_reg.async_get_or_create( + entity_registry.async_get_or_create( DOMAIN, "test", "5678", diff --git a/script/scaffold/templates/device_condition/tests/test_device_condition.py b/script/scaffold/templates/device_condition/tests/test_device_condition.py index 539a60ded97..59781e24fe6 100644 --- a/script/scaffold/templates/device_condition/tests/test_device_condition.py +++ b/script/scaffold/templates/device_condition/tests/test_device_condition.py @@ -8,7 +8,7 @@ from homeassistant.components.NEW_DOMAIN import DOMAIN from homeassistant.components.device_automation import DeviceAutomationType from homeassistant.const import STATE_OFF, STATE_ON from homeassistant.core import HomeAssistant, ServiceCall -from homeassistant.helpers import device_registry, entity_registry +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.setup import async_setup_component from tests.common import ( @@ -16,23 +16,9 @@ from tests.common import ( assert_lists_same, async_get_device_automations, async_mock_service, - mock_device_registry, - mock_registry, ) -@pytest.fixture -def device_reg(hass: HomeAssistant) -> device_registry.DeviceRegistry: - """Return an empty, loaded, registry.""" - return mock_device_registry(hass) - - -@pytest.fixture -def entity_reg(hass: HomeAssistant) -> entity_registry.EntityRegistry: - """Return an empty, loaded, registry.""" - return mock_registry(hass) - - @pytest.fixture def calls(hass: HomeAssistant) -> list[ServiceCall]: """Track calls to a mock service.""" @@ -41,17 +27,19 @@ def calls(hass: HomeAssistant) -> list[ServiceCall]: async def test_get_conditions( hass: HomeAssistant, - device_reg: device_registry.DeviceRegistry, - entity_reg: entity_registry.EntityRegistry, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, ) -> None: """Test we get the expected conditions from a NEW_DOMAIN.""" config_entry = MockConfigEntry(domain="test", data={}) config_entry.add_to_hass(hass) - device_entry = device_reg.async_get_or_create( + device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + entity_registry.async_get_or_create( + DOMAIN, "test", "5678", device_id=device_entry.id ) - entity_reg.async_get_or_create(DOMAIN, "test", "5678", device_id=device_entry.id) expected_conditions = [ { "condition": "device", diff --git a/script/scaffold/templates/device_trigger/tests/test_device_trigger.py b/script/scaffold/templates/device_trigger/tests/test_device_trigger.py index 59ba9654566..8cfaa6f64a9 100644 --- a/script/scaffold/templates/device_trigger/tests/test_device_trigger.py +++ b/script/scaffold/templates/device_trigger/tests/test_device_trigger.py @@ -5,7 +5,8 @@ from homeassistant.components import automation from homeassistant.components.NEW_DOMAIN import DOMAIN from homeassistant.components.device_automation import DeviceAutomationType from homeassistant.const import STATE_OFF, STATE_ON -from homeassistant.helpers import device_registry +from homeassistant.core import HomeAssistant, ServiceCall +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.setup import async_setup_component from tests.common import ( @@ -13,38 +14,30 @@ from tests.common import ( assert_lists_same, async_get_device_automations, async_mock_service, - mock_device_registry, - mock_registry, ) @pytest.fixture -def device_reg(hass): - """Return an empty, loaded, registry.""" - return mock_device_registry(hass) - - -@pytest.fixture -def entity_reg(hass): - """Return an empty, loaded, registry.""" - return mock_registry(hass) - - -@pytest.fixture -def calls(hass): +def calls(hass: HomeAssistant) -> list[ServiceCall]: """Track calls to a mock service.""" return async_mock_service(hass, "test", "automation") -async def test_get_triggers(hass, device_reg, entity_reg): +async def test_get_triggers( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, +) -> None: """Test we get the expected triggers from a NEW_DOMAIN.""" config_entry = MockConfigEntry(domain="test", data={}) config_entry.add_to_hass(hass) - device_entry = device_reg.async_get_or_create( + device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, - connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + entity_registry.async_get_or_create( + DOMAIN, "test", "5678", device_id=device_entry.id ) - entity_reg.async_get_or_create(DOMAIN, "test", "5678", device_id=device_entry.id) expected_triggers = [ { "platform": "device", @@ -67,7 +60,9 @@ async def test_get_triggers(hass, device_reg, entity_reg): assert_lists_same(triggers, expected_triggers) -async def test_if_fires_on_state_change(hass, calls): +async def test_if_fires_on_state_change( + hass: HomeAssistant, calls: list[ServiceCall] +) -> None: """Test for turn_on and turn_off triggers firing.""" hass.states.async_set("NEW_DOMAIN.entity", STATE_OFF)