NextDNS tests improvements (#150791)

tests/components/nextdns/__init__.py

@@ -13,8 +13,6 @@ from nextdns import (
     Settings,
 )
 
-from homeassistant.components.nextdns.const import CONF_PROFILE_ID, DOMAIN
-from homeassistant.const import CONF_API_KEY
 from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry
@@ -155,20 +153,12 @@ def mock_nextdns():
         yield
 
 
-async def init_integration(hass: HomeAssistant) -> MockConfigEntry:
+async def init_integration(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Set up the NextDNS integration in Home Assistant."""
-    entry = MockConfigEntry(
+    mock_config_entry.add_to_hass(hass)
-        domain=DOMAIN,
-        title="Fake Profile",
-        unique_id="xyz12",
-        data={CONF_API_KEY: "fake_api_key", CONF_PROFILE_ID: "xyz12"},
-        entry_id="d9aa37407ddac7b964a99e86312288d6",
-    )
-
-    entry.add_to_hass(hass)
 
     with mock_nextdns():
-        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
         await hass.async_block_till_done()
-
-    return entry

tests/components/nextdns/conftest.py

@@ -0,0 +1,32 @@
+"""Common fixtures for the NextDNS tests."""
+
+from collections.abc import Generator
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from homeassistant.components.nextdns.const import CONF_PROFILE_ID, DOMAIN
+from homeassistant.const import CONF_API_KEY
+
+from tests.common import MockConfigEntry
+
+
+@pytest.fixture
+def mock_setup_entry() -> Generator[AsyncMock]:
+    """Override async_setup_entry."""
+    with patch(
+        "homeassistant.components.nextdns.async_setup_entry", return_value=True
+    ) as mock_setup_entry:
+        yield mock_setup_entry
+
+
+@pytest.fixture
+def mock_config_entry() -> MockConfigEntry:
+    """Return the default mocked config entry."""
+    return MockConfigEntry(
+        domain=DOMAIN,
+        title="Fake Profile",
+        unique_id="xyz12",
+        data={CONF_API_KEY: "fake_api_key", CONF_PROFILE_ID: "xyz12"},
+        entry_id="d9aa37407ddac7b964a99e86312288d6",
+    )

tests/components/nextdns/test_binary_sensor.py

@@ -3,56 +3,65 @@
 from datetime import timedelta
 from unittest.mock import patch
 
+from freezegun.api import FrozenDateTimeFactory
 from nextdns import ApiError
 from syrupy.assertion import SnapshotAssertion
 
-from homeassistant.const import STATE_ON, STATE_UNAVAILABLE, Platform
+from homeassistant.const import STATE_UNAVAILABLE, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
-from homeassistant.util.dt import utcnow
 
 from . import init_integration, mock_nextdns
 
-from tests.common import async_fire_time_changed, snapshot_platform
+from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 
 
 async def test_binary_sensor(
-    hass: HomeAssistant, entity_registry: er.EntityRegistry, snapshot: SnapshotAssertion
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    snapshot: SnapshotAssertion,
+    mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test states of the binary sensors."""
     with patch("homeassistant.components.nextdns.PLATFORMS", [Platform.BINARY_SENSOR]):
-        entry = await init_integration(hass)
+        await init_integration(hass, mock_config_entry)
 
-    await snapshot_platform(hass, entity_registry, snapshot, entry.entry_id)
+    await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)
 
 
-async def test_availability(hass: HomeAssistant) -> None:
+async def test_availability(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+) -> None:
     """Ensure that we mark the entities unavailable correctly when service causes an error."""
-    await init_integration(hass)
+    with patch("homeassistant.components.nextdns.PLATFORMS", [Platform.BINARY_SENSOR]):
+        await init_integration(hass, mock_config_entry)
 
-    state = hass.states.get("binary_sensor.fake_profile_device_connection_status")
+    entity_entries = er.async_entries_for_config_entry(
-    assert state
+        entity_registry, mock_config_entry.entry_id
-    assert state.state != STATE_UNAVAILABLE
+    )
-    assert state.state == STATE_ON
+    entity_ids = (entry.entity_id for entry in entity_entries)
 
-    future = utcnow() + timedelta(minutes=10)
+    for entity_id in entity_ids:
+        assert hass.states.get(entity_id).state != STATE_UNAVAILABLE
+
+    freezer.tick(timedelta(minutes=10))
     with patch(
         "homeassistant.components.nextdns.NextDns.connection_status",
         side_effect=ApiError("API Error"),
     ):
-        async_fire_time_changed(hass, future)
+        async_fire_time_changed(hass)
         await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get("binary_sensor.fake_profile_device_connection_status")
+    for entity_id in entity_ids:
-    assert state
+        assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
-    assert state.state == STATE_UNAVAILABLE
 
-    future = utcnow() + timedelta(minutes=20)
+    freezer.tick(timedelta(minutes=10))
     with mock_nextdns():
-        async_fire_time_changed(hass, future)
+        async_fire_time_changed(hass)
         await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get("binary_sensor.fake_profile_device_connection_status")
+    for entity_id in entity_ids:
-    assert state
+        assert hass.states.get(entity_id).state != STATE_UNAVAILABLE
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == STATE_ON

tests/components/nextdns/test_button.py

@@ -15,31 +15,34 @@ from homeassistant.const import ATTR_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
-from homeassistant.util import dt as dt_util
 
 from . import init_integration
 
-from tests.common import snapshot_platform
+from tests.common import MockConfigEntry, snapshot_platform
 
 
 async def test_button(
-    hass: HomeAssistant, entity_registry: er.EntityRegistry, snapshot: SnapshotAssertion
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    snapshot: SnapshotAssertion,
+    mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test states of the button."""
     with patch("homeassistant.components.nextdns.PLATFORMS", [Platform.BUTTON]):
-        entry = await init_integration(hass)
+        await init_integration(hass, mock_config_entry)
 
-    await snapshot_platform(hass, entity_registry, snapshot, entry.entry_id)
+    await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)
 
 
-async def test_button_press(hass: HomeAssistant) -> None:
+@pytest.mark.freeze_time("2023-10-21")
+async def test_button_press(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test button press."""
-    await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
-    now = dt_util.utcnow()
     with (
         patch("homeassistant.components.nextdns.NextDns.clear_logs") as mock_clear_logs,
-        patch("homeassistant.core.dt_util.utcnow", return_value=now),
     ):
         await hass.services.async_call(
             BUTTON_DOMAIN,
@@ -53,7 +56,7 @@ async def test_button_press(hass: HomeAssistant) -> None:
 
     state = hass.states.get("button.fake_profile_clear_logs")
     assert state
-    assert state.state == now.isoformat()
+    assert state.state == "2023-10-21T00:00:00+00:00"
 
 
 @pytest.mark.parametrize(
@@ -65,9 +68,11 @@ async def test_button_press(hass: HomeAssistant) -> None:
         ClientError,
     ],
 )
-async def test_button_failure(hass: HomeAssistant, exc: Exception) -> None:
+async def test_button_failure(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry, exc: Exception
+) -> None:
     """Tests that the press action throws HomeAssistantError."""
-    await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     with (
         patch("homeassistant.components.nextdns.NextDns.clear_logs", side_effect=exc),
@@ -84,9 +89,11 @@ async def test_button_failure(hass: HomeAssistant, exc: Exception) -> None:
         )
 
 
-async def test_button_auth_error(hass: HomeAssistant) -> None:
+async def test_button_auth_error(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Tests that the press action starts re-auth flow."""
-    entry = await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     with patch(
         "homeassistant.components.nextdns.NextDns.clear_logs",
@@ -99,7 +106,7 @@ async def test_button_auth_error(hass: HomeAssistant) -> None:
             blocking=True,
         )
 
-    assert entry.state is ConfigEntryState.LOADED
+    assert mock_config_entry.state is ConfigEntryState.LOADED
 
     flows = hass.config_entries.flow.async_progress()
     assert len(flows) == 1
@@ -110,4 +117,4 @@ async def test_button_auth_error(hass: HomeAssistant) -> None:
 
     assert "context" in flow
     assert flow["context"].get("source") == SOURCE_REAUTH
-    assert flow["context"].get("entry_id") == entry.entry_id
+    assert flow["context"].get("entry_id") == mock_config_entry.entry_id

tests/components/nextdns/test_config_flow.py

@@ -1,6 +1,6 @@
 """Define tests for the NextDNS config flow."""
 
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 from nextdns import ApiError, InvalidApiKeyError
 import pytest
@@ -14,8 +14,12 @@ from homeassistant.data_entry_flow import FlowResultType
 
 from . import PROFILES, init_integration, mock_nextdns
 
+from tests.common import MockConfigEntry
 
-async def test_form_create_entry(hass: HomeAssistant) -> None:
+
+async def test_form_create_entry(
+    hass: HomeAssistant, mock_setup_entry: AsyncMock
+) -> None:
     """Test that the user step works."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": SOURCE_USER}
@@ -24,14 +28,9 @@ async def test_form_create_entry(hass: HomeAssistant) -> None:
     assert result["step_id"] == "user"
     assert result["errors"] == {}
 
-    with (
+    with patch(
-        patch(
+        "homeassistant.components.nextdns.NextDns.get_profiles",
-            "homeassistant.components.nextdns.NextDns.get_profiles",
+        return_value=PROFILES,
-            return_value=PROFILES,
-        ),
-        patch(
-            "homeassistant.components.nextdns.async_setup_entry", return_value=True
-        ) as mock_setup_entry,
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -44,12 +43,12 @@ async def test_form_create_entry(hass: HomeAssistant) -> None:
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"], {CONF_PROFILE_NAME: "Fake Profile"}
         )
-        await hass.async_block_till_done()
 
     assert result["type"] is FlowResultType.CREATE_ENTRY
     assert result["title"] == "Fake Profile"
     assert result["data"][CONF_API_KEY] == "fake_api_key"
     assert result["data"][CONF_PROFILE_ID] == "xyz12"
+    assert result["result"].unique_id == "xyz12"
     assert len(mock_setup_entry.mock_calls) == 1
 
 
@@ -64,24 +63,55 @@ async def test_form_create_entry(hass: HomeAssistant) -> None:
     ],
 )
 async def test_form_errors(
-    hass: HomeAssistant, exc: Exception, base_error: str
+    hass: HomeAssistant, mock_setup_entry: AsyncMock, exc: Exception, base_error: str
 ) -> None:
     """Test we handle errors."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": SOURCE_USER}
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["errors"] == {}
+
     with patch(
         "homeassistant.components.nextdns.NextDns.get_profiles", side_effect=exc
     ):
-        result = await hass.config_entries.flow.async_init(
+        result = await hass.config_entries.flow.async_configure(
-            DOMAIN,
+            result["flow_id"],
-            context={"source": SOURCE_USER},
+            {CONF_API_KEY: "fake_api_key"},
-            data={CONF_API_KEY: "fake_api_key"},
         )
 
+    assert result["type"] is FlowResultType.FORM
     assert result["errors"] == {"base": base_error}
 
+    with patch(
+        "homeassistant.components.nextdns.NextDns.get_profiles",
+        return_value=PROFILES,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_API_KEY: "fake_api_key"},
+        )
 
-async def test_form_already_configured(hass: HomeAssistant) -> None:
+        assert result["type"] is FlowResultType.FORM
+        assert result["step_id"] == "profiles"
+
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {CONF_PROFILE_NAME: "Fake Profile"}
+        )
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["title"] == "Fake Profile"
+    assert result["data"][CONF_API_KEY] == "fake_api_key"
+    assert result["data"][CONF_PROFILE_ID] == "xyz12"
+    assert result["result"].unique_id == "xyz12"
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_form_already_configured(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test that errors are shown when duplicates are added."""
-    await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": SOURCE_USER}
@@ -103,11 +133,13 @@ async def test_form_already_configured(hass: HomeAssistant) -> None:
     assert result["reason"] == "already_configured"
 
 
-async def test_reauth_successful(hass: HomeAssistant) -> None:
+async def test_reauth_successful(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test starting a reauthentication flow."""
-    entry = await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
-    result = await entry.start_reauth_flow(hass)
+    result = await mock_config_entry.start_reauth_flow(hass)
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "reauth_confirm"
 
@@ -122,7 +154,6 @@ async def test_reauth_successful(hass: HomeAssistant) -> None:
             result["flow_id"],
             user_input={CONF_API_KEY: "new_api_key"},
         )
-        await hass.async_block_till_done()
 
     assert result["type"] is FlowResultType.ABORT
     assert result["reason"] == "reauth_successful"
@@ -139,12 +170,15 @@ async def test_reauth_successful(hass: HomeAssistant) -> None:
     ],
 )
 async def test_reauth_errors(
-    hass: HomeAssistant, exc: Exception, base_error: str
+    hass: HomeAssistant,
+    exc: Exception,
+    base_error: str,
+    mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test reauthentication flow with errors."""
-    entry = await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
-    result = await entry.start_reauth_flow(hass)
+    result = await mock_config_entry.start_reauth_flow(hass)
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "reauth_confirm"
 
@@ -155,6 +189,20 @@ async def test_reauth_errors(
             result["flow_id"],
             user_input={CONF_API_KEY: "new_api_key"},
         )
-        await hass.async_block_till_done()
 
     assert result["errors"] == {"base": base_error}
+
+    with (
+        patch(
+            "homeassistant.components.nextdns.NextDns.get_profiles",
+            return_value=PROFILES,
+        ),
+        mock_nextdns(),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            user_input={CONF_API_KEY: "new_api_key"},
+        )
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "reauth_successful"

tests/components/nextdns/test_coordinator.py

@@ -12,17 +12,18 @@ from homeassistant.core import HomeAssistant
 
 from . import init_integration
 
-from tests.common import async_fire_time_changed
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 async def test_auth_error(
     hass: HomeAssistant,
     freezer: FrozenDateTimeFactory,
+    mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test authentication error when polling data."""
-    entry = await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
-    assert entry.state is ConfigEntryState.LOADED
+    assert mock_config_entry.state is ConfigEntryState.LOADED
 
     freezer.tick(timedelta(minutes=10))
     with (
@@ -62,7 +63,7 @@ async def test_auth_error(
         async_fire_time_changed(hass)
         await hass.async_block_till_done()
 
-    assert entry.state is ConfigEntryState.LOADED
+    assert mock_config_entry.state is ConfigEntryState.LOADED
 
     flows = hass.config_entries.flow.async_progress()
     assert len(flows) == 1
@@ -73,4 +74,4 @@ async def test_auth_error(
 
     assert "context" in flow
     assert flow["context"].get("source") == SOURCE_REAUTH
-    assert flow["context"].get("entry_id") == entry.entry_id
+    assert flow["context"].get("entry_id") == mock_config_entry.entry_id

tests/components/nextdns/test_diagnostics.py

@@ -7,6 +7,7 @@ from homeassistant.core import HomeAssistant
 
 from . import init_integration
 
+from tests.common import MockConfigEntry
 from tests.components.diagnostics import get_diagnostics_for_config_entry
 from tests.typing import ClientSessionGenerator
 
@@ -15,10 +16,11 @@ async def test_entry_diagnostics(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,
     snapshot: SnapshotAssertion,
+    mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test config entry diagnostics."""
-    entry = await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
-    assert await get_diagnostics_for_config_entry(hass, hass_client, entry) == snapshot(
+    assert await get_diagnostics_for_config_entry(
-        exclude=props("created_at", "modified_at")
+        hass, hass_client, mock_config_entry
-    )
+    ) == snapshot(exclude=props("created_at", "modified_at"))

tests/components/nextdns/test_init.py

@@ -6,9 +6,9 @@ from nextdns import ApiError, InvalidApiKeyError
 import pytest
 from tenacity import RetryError
 
-from homeassistant.components.nextdns.const import CONF_PROFILE_ID, DOMAIN
+from homeassistant.components.nextdns.const import DOMAIN
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
-from homeassistant.const import CONF_API_KEY, STATE_UNAVAILABLE
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 
 from . import init_integration
@@ -16,9 +16,11 @@ from . import init_integration
 from tests.common import MockConfigEntry
 
 
-async def test_async_setup_entry(hass: HomeAssistant) -> None:
+async def test_async_setup_entry(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test a successful setup entry."""
-    await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     state = hass.states.get("sensor.fake_profile_dns_queries_blocked_ratio")
     assert state is not None
@@ -29,55 +31,48 @@ async def test_async_setup_entry(hass: HomeAssistant) -> None:
 @pytest.mark.parametrize(
     "exc", [ApiError("API Error"), RetryError("Retry Error"), TimeoutError]
 )
-async def test_config_not_ready(hass: HomeAssistant, exc: Exception) -> None:
+async def test_config_not_ready(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry, exc: Exception
+) -> None:
     """Test for setup failure if the connection to the service fails."""
-    entry = MockConfigEntry(
-        domain=DOMAIN,
-        title="Fake Profile",
-        unique_id="xyz12",
-        data={CONF_API_KEY: "fake_api_key", CONF_PROFILE_ID: "xyz12"},
-    )
-
     with patch(
         "homeassistant.components.nextdns.NextDns.get_profiles",
         side_effect=exc,
     ):
-        entry.add_to_hass(hass)
+        mock_config_entry.add_to_hass(hass)
-        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
-        assert entry.state is ConfigEntryState.SETUP_RETRY
+        assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
 
 
-async def test_unload_entry(hass: HomeAssistant) -> None:
+async def test_unload_entry(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test successful unload of entry."""
-    entry = await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     assert len(hass.config_entries.async_entries(DOMAIN)) == 1
-    assert entry.state is ConfigEntryState.LOADED
+    assert mock_config_entry.state is ConfigEntryState.LOADED
 
-    assert await hass.config_entries.async_unload(entry.entry_id)
+    assert await hass.config_entries.async_unload(mock_config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert entry.state is ConfigEntryState.NOT_LOADED
+    assert mock_config_entry.state is ConfigEntryState.NOT_LOADED
     assert not hass.data.get(DOMAIN)
 
 
-async def test_config_auth_failed(hass: HomeAssistant) -> None:
+async def test_config_auth_failed(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test for setup failure if the auth fails."""
-    entry = MockConfigEntry(
+    mock_config_entry.add_to_hass(hass)
-        domain=DOMAIN,
-        title="Fake Profile",
-        unique_id="xyz12",
-        data={CONF_API_KEY: "fake_api_key", CONF_PROFILE_ID: "xyz12"},
-    )
-    entry.add_to_hass(hass)
 
     with patch(
         "homeassistant.components.nextdns.NextDns.get_profiles",
         side_effect=InvalidApiKeyError,
     ):
-        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
 
-    assert entry.state is ConfigEntryState.SETUP_ERROR
+    assert mock_config_entry.state is ConfigEntryState.SETUP_ERROR
 
     flows = hass.config_entries.flow.async_progress()
     assert len(flows) == 1
@@ -88,4 +83,4 @@ async def test_config_auth_failed(hass: HomeAssistant) -> None:
 
     assert "context" in flow
     assert flow["context"].get("source") == SOURCE_REAUTH
-    assert flow["context"].get("entry_id") == entry.entry_id
+    assert flow["context"].get("entry_id") == mock_config_entry.entry_id

tests/components/nextdns/test_sensor.py

@@ -3,6 +3,7 @@
 from datetime import timedelta
 from unittest.mock import patch
 
+from freezegun.api import FrozenDateTimeFactory
 from nextdns import ApiError
 import pytest
 from syrupy.assertion import SnapshotAssertion
@@ -10,11 +11,10 @@ from syrupy.assertion import SnapshotAssertion
 from homeassistant.const import STATE_UNAVAILABLE, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
-from homeassistant.util.dt import utcnow
 
 from . import init_integration, mock_nextdns
 
-from tests.common import async_fire_time_changed, snapshot_platform
+from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -22,48 +22,35 @@ async def test_sensor(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     snapshot: SnapshotAssertion,
+    mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test states of sensors."""
     with patch("homeassistant.components.nextdns.PLATFORMS", [Platform.SENSOR]):
-        entry = await init_integration(hass)
+        await init_integration(hass, mock_config_entry)
 
-    await snapshot_platform(hass, entity_registry, snapshot, entry.entry_id)
+    await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 async def test_availability(
     hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    mock_config_entry: MockConfigEntry,
     entity_registry: er.EntityRegistry,
 ) -> None:
     """Ensure that we mark the entities unavailable correctly when service causes an error."""
-    await init_integration(hass)
+    with patch("homeassistant.components.nextdns.PLATFORMS", [Platform.SENSOR]):
+        await init_integration(hass, mock_config_entry)
 
-    state = hass.states.get("sensor.fake_profile_dns_queries")
+    entity_entries = er.async_entries_for_config_entry(
-    assert state
+        entity_registry, mock_config_entry.entry_id
-    assert state.state != STATE_UNAVAILABLE
+    )
-    assert state.state == "100"
+    entity_ids = (entry.entity_id for entry in entity_entries)
 
-    state = hass.states.get("sensor.fake_profile_dns_over_https_queries")
+    for entity_id in entity_ids:
-    assert state
+        assert hass.states.get(entity_id).state != STATE_UNAVAILABLE
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "20"
 
-    state = hass.states.get("sensor.fake_profile_dnssec_validated_queries")
+    freezer.tick(timedelta(minutes=10))
-    assert state
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "75"
-
-    state = hass.states.get("sensor.fake_profile_encrypted_queries")
-    assert state
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "60"
-
-    state = hass.states.get("sensor.fake_profile_ipv4_queries")
-    assert state
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "90"
-
-    future = utcnow() + timedelta(minutes=10)
     with (
         patch(
             "homeassistant.components.nextdns.NextDns.get_analytics_status",
@@ -86,55 +73,16 @@ async def test_availability(
             side_effect=ApiError("API Error"),
         ),
     ):
-        async_fire_time_changed(hass, future)
+        async_fire_time_changed(hass)
         await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get("sensor.fake_profile_dns_queries")
+    for entity_id in entity_ids:
-    assert state
+        assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
-    assert state.state == STATE_UNAVAILABLE
 
-    state = hass.states.get("sensor.fake_profile_dns_over_https_queries")
+    freezer.tick(timedelta(minutes=10))
-    assert state
-    assert state.state == STATE_UNAVAILABLE
-
-    state = hass.states.get("sensor.fake_profile_dnssec_validated_queries")
-    assert state
-    assert state.state == STATE_UNAVAILABLE
-
-    state = hass.states.get("sensor.fake_profile_encrypted_queries")
-    assert state
-    assert state.state == STATE_UNAVAILABLE
-
-    state = hass.states.get("sensor.fake_profile_ipv4_queries")
-    assert state
-    assert state.state == STATE_UNAVAILABLE
-
-    future = utcnow() + timedelta(minutes=20)
     with mock_nextdns():
-        async_fire_time_changed(hass, future)
+        async_fire_time_changed(hass)
         await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get("sensor.fake_profile_dns_queries")
+    for entity_id in entity_ids:
-    assert state
+        assert hass.states.get(entity_id).state != STATE_UNAVAILABLE
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "100"
-
-    state = hass.states.get("sensor.fake_profile_dns_over_https_queries")
-    assert state
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "20"
-
-    state = hass.states.get("sensor.fake_profile_dnssec_validated_queries")
-    assert state
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "75"
-
-    state = hass.states.get("sensor.fake_profile_encrypted_queries")
-    assert state
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "60"
-
-    state = hass.states.get("sensor.fake_profile_ipv4_queries")
-    assert state
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == "90"

tests/components/nextdns/test_switch.py

@@ -5,6 +5,7 @@ from unittest.mock import Mock, patch
 
 from aiohttp import ClientError
 from aiohttp.client_exceptions import ClientConnectorError
+from freezegun.api import FrozenDateTimeFactory
 from nextdns import ApiError, InvalidApiKeyError
 import pytest
 from syrupy.assertion import SnapshotAssertion
@@ -25,11 +26,10 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
-from homeassistant.util.dt import utcnow
 
 from . import init_integration, mock_nextdns
 
-from tests.common import async_fire_time_changed, snapshot_platform
+from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
@@ -37,17 +37,20 @@ async def test_switch(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     snapshot: SnapshotAssertion,
+    mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test states of the switches."""
     with patch("homeassistant.components.nextdns.PLATFORMS", [Platform.SWITCH]):
-        entry = await init_integration(hass)
+        await init_integration(hass, mock_config_entry)
 
-    await snapshot_platform(hass, entity_registry, snapshot, entry.entry_id)
+    await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)
 
 
-async def test_switch_on(hass: HomeAssistant) -> None:
+async def test_switch_on(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test the switch can be turned on."""
-    await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     state = hass.states.get("switch.fake_profile_block_page")
     assert state
@@ -71,9 +74,11 @@ async def test_switch_on(hass: HomeAssistant) -> None:
         mock_switch_on.assert_called_once()
 
 
-async def test_switch_off(hass: HomeAssistant) -> None:
+async def test_switch_off(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Test the switch can be turned on."""
-    await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     state = hass.states.get("switch.fake_profile_web3")
     assert state
@@ -97,6 +102,7 @@ async def test_switch_off(hass: HomeAssistant) -> None:
         mock_switch_on.assert_called_once()
 
 
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
 @pytest.mark.parametrize(
     "exc",
     [
@@ -105,36 +111,43 @@ async def test_switch_off(hass: HomeAssistant) -> None:
         TimeoutError,
     ],
 )
-async def test_availability(hass: HomeAssistant, exc: Exception) -> None:
+async def test_availability(
+    hass: HomeAssistant,
+    exc: Exception,
+    freezer: FrozenDateTimeFactory,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+) -> None:
     """Ensure that we mark the entities unavailable correctly when service causes an error."""
-    await init_integration(hass)
+    with patch("homeassistant.components.nextdns.PLATFORMS", [Platform.SWITCH]):
+        await init_integration(hass, mock_config_entry)
 
-    state = hass.states.get("switch.fake_profile_web3")
+    entity_entries = er.async_entries_for_config_entry(
-    assert state
+        entity_registry, mock_config_entry.entry_id
-    assert state.state != STATE_UNAVAILABLE
+    )
-    assert state.state == STATE_ON
+    entity_ids = (entry.entity_id for entry in entity_entries)
 
-    future = utcnow() + timedelta(minutes=10)
+    for entity_id in entity_ids:
+        assert hass.states.get(entity_id).state != STATE_UNAVAILABLE
+
+    freezer.tick(timedelta(minutes=10))
     with patch(
         "homeassistant.components.nextdns.NextDns.get_settings",
         side_effect=exc,
     ):
-        async_fire_time_changed(hass, future)
+        async_fire_time_changed(hass)
         await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get("switch.fake_profile_web3")
+    for entity_id in entity_ids:
-    assert state
+        assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
-    assert state.state == STATE_UNAVAILABLE
 
-    future = utcnow() + timedelta(minutes=20)
+    freezer.tick(timedelta(minutes=10))
     with mock_nextdns():
-        async_fire_time_changed(hass, future)
+        async_fire_time_changed(hass)
         await hass.async_block_till_done(wait_background_tasks=True)
 
-    state = hass.states.get("switch.fake_profile_web3")
+    for entity_id in entity_ids:
-    assert state
+        assert hass.states.get(entity_id).state != STATE_UNAVAILABLE
-    assert state.state != STATE_UNAVAILABLE
-    assert state.state == STATE_ON
 
 
 @pytest.mark.parametrize(
@@ -146,9 +159,11 @@ async def test_availability(hass: HomeAssistant, exc: Exception) -> None:
         ClientError,
     ],
 )
-async def test_switch_failure(hass: HomeAssistant, exc: Exception) -> None:
+async def test_switch_failure(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry, exc: Exception
+) -> None:
     """Tests that the turn on/off service throws HomeAssistantError."""
-    await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     with (
         patch("homeassistant.components.nextdns.NextDns.set_setting", side_effect=exc),
@@ -162,9 +177,11 @@ async def test_switch_failure(hass: HomeAssistant, exc: Exception) -> None:
         )
 
 
-async def test_switch_auth_error(hass: HomeAssistant) -> None:
+async def test_switch_auth_error(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
     """Tests that the turn on/off action starts re-auth flow."""
-    entry = await init_integration(hass)
+    await init_integration(hass, mock_config_entry)
 
     with patch(
         "homeassistant.components.nextdns.NextDns.set_setting",
@@ -177,7 +194,7 @@ async def test_switch_auth_error(hass: HomeAssistant) -> None:
             blocking=True,
         )
 
-    assert entry.state is ConfigEntryState.LOADED
+    assert mock_config_entry.state is ConfigEntryState.LOADED
 
     flows = hass.config_entries.flow.async_progress()
     assert len(flows) == 1
@@ -188,4 +205,4 @@ async def test_switch_auth_error(hass: HomeAssistant) -> None:
 
     assert "context" in flow
     assert flow["context"].get("source") == SOURCE_REAUTH
-    assert flow["context"].get("entry_id") == entry.entry_id
+    assert flow["context"].get("entry_id") == mock_config_entry.entry_id