Re-implement Cloudflare using coordinator (#156817)

Signed-off-by: David Rapan <david@rapan.cz>

homeassistant/components/cloudflare/__init__.py

@@ -2,86 +2,23 @@
 
 from __future__ import annotations
 
-import asyncio
-from dataclasses import dataclass
-from datetime import datetime, timedelta
-import logging
-import socket
-
-import pycfdns
-
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_API_TOKEN, CONF_ZONE
 from homeassistant.core import HomeAssistant, ServiceCall
-from homeassistant.exceptions import (
-    ConfigEntryAuthFailed,
-    ConfigEntryNotReady,
-    HomeAssistantError,
-)
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.event import async_track_time_interval
-from homeassistant.util.location import async_detect_location_info
-from homeassistant.util.network import is_ipv4_address
 
-from .const import CONF_RECORDS, DEFAULT_UPDATE_INTERVAL, DOMAIN, SERVICE_UPDATE_RECORDS
-
-_LOGGER = logging.getLogger(__name__)
-
-type CloudflareConfigEntry = ConfigEntry[CloudflareRuntimeData]
-
-
-@dataclass
-class CloudflareRuntimeData:
-    """Runtime data for Cloudflare config entry."""
-
-    client: pycfdns.Client
-    dns_zone: pycfdns.ZoneModel
+from .const import DOMAIN, SERVICE_UPDATE_RECORDS
+from .coordinator import CloudflareConfigEntry, CloudflareCoordinator
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: CloudflareConfigEntry) -> bool:
     """Set up Cloudflare from a config entry."""
-    session = async_get_clientsession(hass)
-    client = pycfdns.Client(
-        api_token=entry.data[CONF_API_TOKEN],
-        client_session=session,
-    )
+    entry.runtime_data = CloudflareCoordinator(hass, entry)
+    await entry.runtime_data.async_config_entry_first_refresh()
 
-    try:
-        dns_zones = await client.list_zones()
-        dns_zone = next(
-            zone for zone in dns_zones if zone["name"] == entry.data[CONF_ZONE]
-        )
-    except pycfdns.AuthenticationException as error:
-        raise ConfigEntryAuthFailed from error
-    except pycfdns.ComunicationException as error:
-        raise ConfigEntryNotReady from error
+    # Since we are not using coordinator for data reads, we need to add dummy listener
+    entry.async_on_unload(entry.runtime_data.async_add_listener(lambda: None))
 
-    entry.runtime_data = CloudflareRuntimeData(client, dns_zone)
-
-    async def update_records(now: datetime) -> None:
-        """Set up recurring update."""
-        try:
-            await _async_update_cloudflare(hass, entry)
-        except (
-            pycfdns.AuthenticationException,
-            pycfdns.ComunicationException,
-        ) as error:
-            _LOGGER.error("Error updating zone %s: %s", entry.data[CONF_ZONE], error)
-
-    async def update_records_service(call: ServiceCall) -> None:
+    async def update_records_service(_: ServiceCall) -> None:
         """Set up service for manual trigger."""
-        try:
-            await _async_update_cloudflare(hass, entry)
-        except (
-            pycfdns.AuthenticationException,
-            pycfdns.ComunicationException,
-        ) as error:
-            _LOGGER.error("Error updating zone %s: %s", entry.data[CONF_ZONE], error)
-
-    update_interval = timedelta(minutes=DEFAULT_UPDATE_INTERVAL)
-    entry.async_on_unload(
-        async_track_time_interval(hass, update_records, update_interval)
-    )
+        await entry.runtime_data.async_request_refresh()
 
     hass.services.async_register(DOMAIN, SERVICE_UPDATE_RECORDS, update_records_service)
 
@@ -92,49 +29,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: CloudflareConfigEntry)
     """Unload Cloudflare config entry."""
 
     return True
-
-
-async def _async_update_cloudflare(
-    hass: HomeAssistant,
-    entry: CloudflareConfigEntry,
-) -> None:
-    client = entry.runtime_data.client
-    dns_zone = entry.runtime_data.dns_zone
-    target_records: list[str] = entry.data[CONF_RECORDS]
-
-    _LOGGER.debug("Starting update for zone %s", dns_zone["name"])
-
-    records = await client.list_dns_records(zone_id=dns_zone["id"], type="A")
-    _LOGGER.debug("Records: %s", records)
-
-    session = async_get_clientsession(hass, family=socket.AF_INET)
-    location_info = await async_detect_location_info(session)
-
-    if not location_info or not is_ipv4_address(location_info.ip):
-        raise HomeAssistantError("Could not get external IPv4 address")
-
-    filtered_records = [
-        record
-        for record in records
-        if record["name"] in target_records and record["content"] != location_info.ip
-    ]
-
-    if len(filtered_records) == 0:
-        _LOGGER.debug("All target records are up to date")
-        return
-
-    await asyncio.gather(
-        *[
-            client.update_dns_record(
-                zone_id=dns_zone["id"],
-                record_id=record["id"],
-                record_content=location_info.ip,
-                record_name=record["name"],
-                record_type=record["type"],
-                record_proxied=record["proxied"],
-            )
-            for record in filtered_records
-        ]
-    )
-
-    _LOGGER.debug("Update for zone %s is complete", dns_zone["name"])

homeassistant/components/cloudflare/coordinator.py

@@ -0,0 +1,116 @@
+"""Contains the Coordinator for updating the IP addresses of your Cloudflare DNS records."""
+
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+from logging import getLogger
+import socket
+
+import pycfdns
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_API_TOKEN, CONF_ZONE
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.util.location import async_detect_location_info
+from homeassistant.util.network import is_ipv4_address
+
+from .const import CONF_RECORDS, DEFAULT_UPDATE_INTERVAL
+
+_LOGGER = getLogger(__name__)
+
+type CloudflareConfigEntry = ConfigEntry[CloudflareCoordinator]
+
+
+class CloudflareCoordinator(DataUpdateCoordinator[None]):
+    """Coordinates records updates."""
+
+    config_entry: CloudflareConfigEntry
+    client: pycfdns.Client
+    zone: pycfdns.ZoneModel
+
+    def __init__(
+        self, hass: HomeAssistant, config_entry: CloudflareConfigEntry
+    ) -> None:
+        """Initialize an coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=config_entry,
+            name=config_entry.title,
+            update_interval=timedelta(minutes=DEFAULT_UPDATE_INTERVAL),
+        )
+
+    async def _async_setup(self) -> None:
+        """Set up the coordinator."""
+        self.client = pycfdns.Client(
+            api_token=self.config_entry.data[CONF_API_TOKEN],
+            client_session=async_get_clientsession(self.hass),
+        )
+
+        try:
+            self.zone = next(
+                zone
+                for zone in await self.client.list_zones()
+                if zone["name"] == self.config_entry.data[CONF_ZONE]
+            )
+        except pycfdns.AuthenticationException as e:
+            raise ConfigEntryAuthFailed from e
+        except pycfdns.ComunicationException as e:
+            raise UpdateFailed("Error communicating with API") from e
+
+    async def _async_update_data(self) -> None:
+        """Update records."""
+        _LOGGER.debug("Starting update for zone %s", self.zone["name"])
+        try:
+            records = await self.client.list_dns_records(
+                zone_id=self.zone["id"], type="A"
+            )
+            _LOGGER.debug("Records: %s", records)
+
+            target_records: list[str] = self.config_entry.data[CONF_RECORDS]
+
+            location_info = await async_detect_location_info(
+                async_get_clientsession(self.hass, family=socket.AF_INET)
+            )
+
+            if not location_info or not is_ipv4_address(location_info.ip):
+                raise UpdateFailed("Could not get external IPv4 address")
+
+            filtered_records = [
+                record
+                for record in records
+                if record["name"] in target_records
+                and record["content"] != location_info.ip
+            ]
+
+            if len(filtered_records) == 0:
+                _LOGGER.debug("All target records are up to date")
+                return
+
+            await asyncio.gather(
+                *[
+                    self.client.update_dns_record(
+                        zone_id=self.zone["id"],
+                        record_id=record["id"],
+                        record_content=location_info.ip,
+                        record_name=record["name"],
+                        record_type=record["type"],
+                        record_proxied=record["proxied"],
+                    )
+                    for record in filtered_records
+                ]
+            )
+
+            _LOGGER.debug("Update for zone %s is complete", self.zone["name"])
+
+        except (
+            pycfdns.AuthenticationException,
+            pycfdns.ComunicationException,
+        ) as e:
+            raise UpdateFailed(
+                f"Error updating zone {self.config_entry.data[CONF_ZONE]}"
+            ) from e

tests/components/cloudflare/conftest.py

@@ -5,15 +5,21 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
+from homeassistant.util.location import LocationInfo
+
 from . import get_mock_client
 
+LOCATION_PATCH_TARGET = (
+    "homeassistant.components.cloudflare.coordinator.async_detect_location_info"
+)
+
 
 @pytest.fixture
 def cfupdate() -> Generator[MagicMock]:
     """Mock the CloudflareUpdater for easier testing."""
     mock_cfupdate = get_mock_client()
     with patch(
-        "homeassistant.components.cloudflare.pycfdns.Client",
+        "homeassistant.components.cloudflare.coordinator.pycfdns.Client",
         return_value=mock_cfupdate,
     ) as mock_api:
         yield mock_api
@@ -28,3 +34,25 @@ def cfupdate_flow() -> Generator[MagicMock]:
         return_value=mock_cfupdate,
     ) as mock_api:
         yield mock_api
+
+
+@pytest.fixture
+def location_info() -> Generator[None]:
+    """Mock the LocationInfo for easier testing."""
+    with patch(
+        LOCATION_PATCH_TARGET,
+        return_value=LocationInfo(
+            "0.0.0.0",
+            "US",
+            "USD",
+            "CA",
+            "California",
+            "San Diego",
+            "92122",
+            "America/Los_Angeles",
+            32.8594,
+            -117.2073,
+            True,
+        ),
+    ):
+        yield

tests/components/cloudflare/test_init.py

@@ -3,6 +3,7 @@
 from datetime import timedelta
 from unittest.mock import MagicMock, patch
 
+from freezegun.api import FrozenDateTimeFactory
 import pycfdns
 import pytest
 
@@ -14,29 +15,13 @@ from homeassistant.components.cloudflare.const import (
 )
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import HomeAssistantError
-from homeassistant.util import dt as dt_util
-from homeassistant.util.location import LocationInfo
 
 from . import ENTRY_CONFIG, init_integration
+from .conftest import LOCATION_PATCH_TARGET
 
 from tests.common import MockConfigEntry, async_fire_time_changed
 
 
-async def test_unload_entry(hass: HomeAssistant, cfupdate: MagicMock) -> None:
-    """Test successful unload of entry."""
-    entry = await init_integration(hass)
-
-    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
-    assert entry.state is ConfigEntryState.LOADED
-
-    assert await hass.config_entries.async_unload(entry.entry_id)
-    await hass.async_block_till_done()
-
-    assert entry.state is ConfigEntryState.NOT_LOADED
-    assert not hass.data.get(DOMAIN)
-
-
 @pytest.mark.parametrize(
     "side_effect",
     [pycfdns.ComunicationException()],
@@ -83,6 +68,21 @@ async def test_async_setup_raises_entry_auth_failed(
     assert flow["context"]["entry_id"] == entry.entry_id
 
 
+@pytest.mark.usefixtures("location_info")
+async def test_unload_entry(hass: HomeAssistant, cfupdate: MagicMock) -> None:
+    """Test successful unload of entry."""
+    entry = await init_integration(hass)
+
+    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
+    assert entry.state is ConfigEntryState.LOADED
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.NOT_LOADED
+
+
+@pytest.mark.usefixtures("location_info")
 async def test_integration_services(
     hass: HomeAssistant, cfupdate: MagicMock, caplog: pytest.LogCaptureFixture
 ) -> None:
@@ -92,36 +92,24 @@ async def test_integration_services(
     entry = await init_integration(hass)
     assert entry.state is ConfigEntryState.LOADED
 
-    with patch(
-        "homeassistant.components.cloudflare.async_detect_location_info",
-        return_value=LocationInfo(
-            "0.0.0.0",
-            "US",
-            "USD",
-            "CA",
-            "California",
-            "San Diego",
-            "92122",
-            "America/Los_Angeles",
-            32.8594,
-            -117.2073,
-            True,
-        ),
-    ):
-        await hass.services.async_call(
-            DOMAIN,
-            SERVICE_UPDATE_RECORDS,
-            {},
-            blocking=True,
-        )
-        await hass.async_block_till_done()
+    assert len(instance.update_dns_record.mock_calls) == 2
+    instance.update_dns_record.reset_mock()
+
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_UPDATE_RECORDS,
+        {},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
 
     assert len(instance.update_dns_record.mock_calls) == 2
     assert "All target records are up to date" not in caplog.text
 
 
+@pytest.mark.usefixtures("location_info")
 async def test_integration_services_with_issue(
-    hass: HomeAssistant, cfupdate: MagicMock
+    hass: HomeAssistant, cfupdate: MagicMock, caplog: pytest.LogCaptureFixture
 ) -> None:
     """Test integration services with issue."""
     instance = cfupdate.return_value
@@ -129,13 +117,10 @@ async def test_integration_services_with_issue(
     entry = await init_integration(hass)
     assert entry.state is ConfigEntryState.LOADED
 
-    with (
-        patch(
-            "homeassistant.components.cloudflare.async_detect_location_info",
-            return_value=None,
-        ),
-        pytest.raises(HomeAssistantError, match="Could not get external IPv4 address"),
-    ):
+    assert len(instance.update_dns_record.mock_calls) == 2
+    instance.update_dns_record.reset_mock()
+
+    with patch(LOCATION_PATCH_TARGET, return_value=None):
         await hass.services.async_call(
             DOMAIN,
             SERVICE_UPDATE_RECORDS,
@@ -144,8 +129,10 @@ async def test_integration_services_with_issue(
         )
 
     instance.update_dns_record.assert_not_called()
+    assert "Could not get external IPv4 address" in caplog.text
 
 
+@pytest.mark.usefixtures("location_info")
 async def test_integration_services_with_nonexisting_record(
     hass: HomeAssistant, cfupdate: MagicMock, caplog: pytest.LogCaptureFixture
 ) -> None:
@@ -157,38 +144,24 @@ async def test_integration_services_with_nonexisting_record(
     )
     assert entry.state is ConfigEntryState.LOADED
 
-    with patch(
-        "homeassistant.components.cloudflare.async_detect_location_info",
-        return_value=LocationInfo(
-            "0.0.0.0",
-            "US",
-            "USD",
-            "CA",
-            "California",
-            "San Diego",
-            "92122",
-            "America/Los_Angeles",
-            32.8594,
-            -117.2073,
-            True,
-        ),
-    ):
-        await hass.services.async_call(
-            DOMAIN,
-            SERVICE_UPDATE_RECORDS,
-            {},
-            blocking=True,
-        )
-        await hass.async_block_till_done()
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_UPDATE_RECORDS,
+        {},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
 
     instance.update_dns_record.assert_not_called()
     assert "All target records are up to date" in caplog.text
 
 
+@pytest.mark.usefixtures("location_info")
 async def test_integration_update_interval(
     hass: HomeAssistant,
     cfupdate: MagicMock,
     caplog: pytest.LogCaptureFixture,
+    freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test integration update interval."""
     instance = cfupdate.return_value
@@ -196,39 +169,23 @@ async def test_integration_update_interval(
     entry = await init_integration(hass)
     assert entry.state is ConfigEntryState.LOADED
 
-    with patch(
-        "homeassistant.components.cloudflare.async_detect_location_info",
-        return_value=LocationInfo(
-            "0.0.0.0",
-            "US",
-            "USD",
-            "CA",
-            "California",
-            "San Diego",
-            "92122",
-            "America/Los_Angeles",
-            32.8594,
-            -117.2073,
-            True,
-        ),
-    ):
-        async_fire_time_changed(
-            hass, dt_util.utcnow() + timedelta(minutes=DEFAULT_UPDATE_INTERVAL)
-        )
-        await hass.async_block_till_done(wait_background_tasks=True)
-        assert len(instance.update_dns_record.mock_calls) == 2
-        assert "All target records are up to date" not in caplog.text
+    freezer.tick(timedelta(minutes=DEFAULT_UPDATE_INTERVAL))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done(wait_background_tasks=True)
+    assert len(instance.list_dns_records.mock_calls) == 2
+    assert len(instance.update_dns_record.mock_calls) == 4
+    assert "All target records are up to date" not in caplog.text
 
-        instance.list_dns_records.side_effect = pycfdns.AuthenticationException()
-        async_fire_time_changed(
-            hass, dt_util.utcnow() + timedelta(minutes=DEFAULT_UPDATE_INTERVAL)
-        )
-        await hass.async_block_till_done(wait_background_tasks=True)
-        assert len(instance.update_dns_record.mock_calls) == 2
+    instance.list_dns_records.side_effect = pycfdns.AuthenticationException()
+    freezer.tick(timedelta(minutes=DEFAULT_UPDATE_INTERVAL))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done(wait_background_tasks=True)
+    assert len(instance.list_dns_records.mock_calls) == 3
+    assert len(instance.update_dns_record.mock_calls) == 4
 
-        instance.list_dns_records.side_effect = pycfdns.ComunicationException()
-        async_fire_time_changed(
-            hass, dt_util.utcnow() + timedelta(minutes=DEFAULT_UPDATE_INTERVAL)
-        )
-        await hass.async_block_till_done(wait_background_tasks=True)
-        assert len(instance.update_dns_record.mock_calls) == 2
+    instance.list_dns_records.side_effect = pycfdns.ComunicationException()
+    freezer.tick(timedelta(minutes=DEFAULT_UPDATE_INTERVAL))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done(wait_background_tasks=True)
+    assert len(instance.list_dns_records.mock_calls) == 4
+    assert len(instance.update_dns_record.mock_calls) == 4