Fix missing mocking in nextdns tests (#114541)

home-assistant · Apr 1, 2024 · 96120b6 · 96120b6
1 parent 429b5d2
commit 96120b6
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 50 deletions.
diff --git a/tests/components/nextdns/__init__.py b/tests/components/nextdns/__init__.py
@@ -1,5 +1,6 @@
 """Tests for the NextDNS integration."""
 
+from contextlib import contextmanager
 from unittest.mock import patch
 
 from nextdns import (
@@ -113,16 +114,9 @@
 )
 
 
-async def init_integration(hass: HomeAssistant) -> MockConfigEntry:
-    """Set up the NextDNS integration in Home Assistant."""
-    entry = MockConfigEntry(
-        domain=DOMAIN,
-        title="Fake Profile",
-        unique_id="xyz12",
-        data={CONF_API_KEY: "fake_api_key", CONF_PROFILE_ID: "xyz12"},
-        entry_id="d9aa37407ddac7b964a99e86312288d6",
-    )
-
+@contextmanager
+def mock_nextdns():
+    """Mock the NextDNS class."""
     with (
         patch(
             "homeassistant.components.nextdns.NextDns.get_profiles",
@@ -157,7 +151,22 @@ async def init_integration(hass: HomeAssistant) -> MockConfigEntry:
             return_value=CONNECTION_STATUS,
         ),
     ):
-        entry.add_to_hass(hass)
+        yield
+
+
+async def init_integration(hass: HomeAssistant) -> MockConfigEntry:
+    """Set up the NextDNS integration in Home Assistant."""
+    entry = MockConfigEntry(
+        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.async_block_till_done()
 

diff --git a/tests/components/nextdns/test_binary_sensor.py b/tests/components/nextdns/test_binary_sensor.py
@@ -10,7 +10,7 @@
 from homeassistant.helpers import entity_registry as er
 from homeassistant.util.dt import utcnow
 
-from . import CONNECTION_STATUS, init_integration
+from . import init_integration, mock_nextdns
 
 from tests.common import async_fire_time_changed
 
@@ -57,19 +57,16 @@ async def test_availability(hass: HomeAssistant) -> None:
         side_effect=ApiError("API Error"),
     ):
         async_fire_time_changed(hass, future)
-        await hass.async_block_till_done()
+        await hass.async_block_till_done(wait_background_tasks=True)
 
     state = hass.states.get("binary_sensor.fake_profile_device_connection_status")
     assert state
     assert state.state == STATE_UNAVAILABLE
 
     future = utcnow() + timedelta(minutes=20)
-    with patch(
-        "homeassistant.components.nextdns.NextDns.connection_status",
-        return_value=CONNECTION_STATUS,
-    ):
+    with mock_nextdns():
         async_fire_time_changed(hass, future)
-        await hass.async_block_till_done()
+        await hass.async_block_till_done(wait_background_tasks=True)
 
     state = hass.states.get("binary_sensor.fake_profile_device_connection_status")
     assert state

diff --git a/tests/components/nextdns/test_sensor.py b/tests/components/nextdns/test_sensor.py
@@ -11,7 +11,7 @@
 from homeassistant.helpers import entity_registry as er
 from homeassistant.util.dt import utcnow
 
-from . import DNSSEC, ENCRYPTION, IP_VERSIONS, PROTOCOLS, STATUS, init_integration
+from . import init_integration, mock_nextdns
 
 from tests.common import async_fire_time_changed
 
@@ -332,7 +332,7 @@ async def test_availability(
         ),
     ):
         async_fire_time_changed(hass, future)
-        await hass.async_block_till_done()
+        await hass.async_block_till_done(wait_background_tasks=True)
 
     state = hass.states.get("sensor.fake_profile_dns_queries")
     assert state
@@ -355,30 +355,9 @@ async def test_availability(
     assert state.state == STATE_UNAVAILABLE
 
     future = utcnow() + timedelta(minutes=20)
-    with (
-        patch(
-            "homeassistant.components.nextdns.NextDns.get_analytics_status",
-            return_value=STATUS,
-        ),
-        patch(
-            "homeassistant.components.nextdns.NextDns.get_analytics_encryption",
-            return_value=ENCRYPTION,
-        ),
-        patch(
-            "homeassistant.components.nextdns.NextDns.get_analytics_dnssec",
-            return_value=DNSSEC,
-        ),
-        patch(
-            "homeassistant.components.nextdns.NextDns.get_analytics_ip_versions",
-            return_value=IP_VERSIONS,
-        ),
-        patch(
-            "homeassistant.components.nextdns.NextDns.get_analytics_protocols",
-            return_value=PROTOCOLS,
-        ),
-    ):
+    with mock_nextdns():
         async_fire_time_changed(hass, future)
-        await hass.async_block_till_done()
+        await hass.async_block_till_done(wait_background_tasks=True)
 
     state = hass.states.get("sensor.fake_profile_dns_queries")
     assert state

diff --git a/tests/components/nextdns/test_switch.py b/tests/components/nextdns/test_switch.py
@@ -22,7 +22,7 @@
 from homeassistant.helpers import entity_registry as er
 from homeassistant.util.dt import utcnow
 
-from . import SETTINGS, init_integration
+from . import init_integration, mock_nextdns
 
 from tests.common import async_fire_time_changed
 
@@ -693,19 +693,16 @@ async def test_availability(hass: HomeAssistant) -> None:
         side_effect=ApiError("API Error"),
     ):
         async_fire_time_changed(hass, future)
-        await hass.async_block_till_done()
+        await hass.async_block_till_done(wait_background_tasks=True)
 
     state = hass.states.get("switch.fake_profile_web3")
     assert state
     assert state.state == STATE_UNAVAILABLE
 
     future = utcnow() + timedelta(minutes=20)
-    with patch(
-        "homeassistant.components.nextdns.NextDns.get_settings",
-        return_value=SETTINGS,
-    ):
+    with mock_nextdns():
         async_fire_time_changed(hass, future)
-        await hass.async_block_till_done()
+        await hass.async_block_till_done(wait_background_tasks=True)
 
     state = hass.states.get("switch.fake_profile_web3")
     assert state