Make IOMeter use SSE (#172364)

homeassistant/components/iometer/__init__.py

@@ -1,11 +1,10 @@
 """The IOmeter integration."""
 
-from iometer import IOmeterClient, IOmeterConnectionError
+from iometer import IOmeterSSEClient
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
 from .coordinator import IOmeterConfigEntry, IOMeterCoordinator
@@ -15,19 +14,21 @@ PLATFORMS: list[Platform] = [Platform.BINARY_SENSOR, Platform.SENSOR]
 
 async def async_setup_entry(hass: HomeAssistant, entry: IOmeterConfigEntry) -> bool:
     """Set up IOmeter from a config entry."""
-
     host = entry.data[CONF_HOST]
     session = async_get_clientsession(hass)
-    client = IOmeterClient(host=host, session=session)
-    try:
-        await client.get_current_status()
-    except IOmeterConnectionError as err:
-        raise ConfigEntryNotReady from err
+    client = IOmeterSSEClient(host=host, session=session)
 
     coordinator = IOMeterCoordinator(hass, entry, client)
-    await coordinator.async_config_entry_first_refresh()
+    await coordinator.async_start()
+    try:
+        await coordinator.async_config_entry_first_refresh()
+    except Exception:
+        await coordinator.async_stop()
+        raise
+
     entry.runtime_data = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+    entry.async_on_unload(coordinator.async_stop)
 
     return True
 

homeassistant/components/iometer/config_flow.py

@@ -5,8 +5,8 @@ from typing import Any, Final, override
 from iometer import (
     IOmeterClient,
     IOmeterConnectionError,
-    IOmeterNoReadingsError,
     IOmeterNoStatusError,
+    IOmeterTimeoutError,
 )
 import voluptuous as vol
 
@@ -40,19 +40,17 @@ class IOMeterConfigFlow(ConfigFlow, domain=DOMAIN):
         client = IOmeterClient(host=host, session=session)
         try:
             status = await client.get_current_status()
-            _ = await client.get_current_reading()
         except IOmeterNoStatusError:
             return self.async_abort(reason="no_status")
-        except IOmeterNoReadingsError:
-            return self.async_abort(reason="no_readings")
-        except IOmeterConnectionError:
+        except IOmeterTimeoutError, IOmeterConnectionError:
             return self.async_abort(reason="cannot_connect")
 
-        self._meter_number = status.meter.number
+        if not status.meter:
+            return self.async_abort(reason="no_readings")
 
+        self._meter_number = status.meter.number
         await self.async_set_unique_id(status.device.id)
         self._abort_if_unique_id_configured()
-
         self.context["title_placeholders"] = {"name": f"IOmeter {self._meter_number}"}
         return await self.async_step_zeroconf_confirm()
 
@@ -82,18 +80,19 @@ class IOMeterConfigFlow(ConfigFlow, domain=DOMAIN):
             client = IOmeterClient(host=self._host, session=session)
             try:
                 status = await client.get_current_status()
-                _ = await client.get_current_reading()
             except IOmeterNoStatusError:
                 errors["base"] = "no_status"
-            except IOmeterNoReadingsError:
-                errors["base"] = "no_readings"
-            except IOmeterConnectionError:
+            except IOmeterTimeoutError, IOmeterConnectionError:
                 errors["base"] = "cannot_connect"
             else:
-                self._meter_number = status.meter.number
-                await self.async_set_unique_id(status.device.id)
-                self._abort_if_unique_id_configured()
-                return await self._async_create_entry()
+                if not status.meter:
+                    errors["base"] = "no_readings"
+                else:
+                    self._meter_number = status.meter.number
+                    await self.async_set_unique_id(status.device.id)
+                    self._abort_if_unique_id_configured()
+                    return await self._async_create_entry()
+
         return self.async_show_form(
             step_id="user",
             data_schema=CONFIG_SCHEMA,

homeassistant/components/iometer/coordinator.py

@@ -1,22 +1,30 @@
 """DataUpdateCoordinator for IOmeter."""
 
+import asyncio
+from collections.abc import Callable
+import contextlib
 from dataclasses import dataclass
-from datetime import timedelta
 import logging
 from typing import override
 
-from iometer import IOmeterClient, IOmeterConnectionError, Reading, Status
+from iometer import (
+    IOmeterConnectionError,
+    IOmeterNoReadingsError,
+    IOmeterNoStatusError,
+    IOmeterSSEClient,
+    IOmeterTimeoutError,
+    Reading,
+    Status,
+)
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
-from homeassistant.helpers.debounce import Debouncer
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
-DEFAULT_SCAN_INTERVAL = timedelta(seconds=10)
 
 type IOmeterConfigEntry = ConfigEntry[IOMeterCoordinator]
 
@@ -33,50 +41,135 @@ class IOMeterCoordinator(DataUpdateCoordinator[IOmeterData]):
     """Class to manage fetching IOmeter data."""
 
     config_entry: IOmeterConfigEntry
-    client: IOmeterClient
+    client: IOmeterSSEClient
     current_fw_version: str = ""
 
     def __init__(
         self,
         hass: HomeAssistant,
         config_entry: IOmeterConfigEntry,
-        client: IOmeterClient,
+        client: IOmeterSSEClient,
     ) -> None:
         """Initialize coordinator."""
-
         super().__init__(
             hass,
             _LOGGER,
             config_entry=config_entry,
             name=DOMAIN,
-            update_interval=DEFAULT_SCAN_INTERVAL,
-            request_refresh_debouncer=Debouncer(
-                hass, _LOGGER, cooldown=1.0, immediate=False
-            ),
         )
         self.client = client
         self.identifier = config_entry.entry_id
+        self._reading: Reading | None = None
+        self._status: Status | None = None
+        self._first_data_event: asyncio.Event = asyncio.Event()
+        self._cancel_readings: Callable[[], None] | None = None
+        self._cancel_status: Callable[[], None] | None = None
+        self._readings_task: asyncio.Task | None = None
+        self._status_task: asyncio.Task | None = None
+
+    async def async_start(self) -> None:
+        """Register SSE subscriptions."""
+        self._cancel_readings = self.client.subscribe_readings(
+            self._on_reading,
+            self._on_reading_error,
+        )
+        self._readings_task = getattr(self._cancel_readings, "__self__", None)
+        self._cancel_status = self.client.subscribe_status(
+            self._on_status,
+            self._on_status_error,
+        )
+        self._status_task = getattr(self._cancel_status, "__self__", None)
+
+    async def async_stop(self) -> None:
+        """Cancel SSE subscriptions and await task teardown."""
+        if self._cancel_readings:
+            self._cancel_readings()
+            self._cancel_readings = None
+        if self._cancel_status:
+            self._cancel_status()
+            self._cancel_status = None
+        for task in (self._readings_task, self._status_task):
+            if task and not task.done():
+                with contextlib.suppress(asyncio.CancelledError):
+                    await task
+        self._readings_task = None
+        self._status_task = None
 
     @override
     async def _async_update_data(self) -> IOmeterData:
-        """Update data async."""
+        """Wait for first SSE data; subsequent updates arrive via async_set_updated_data."""
         try:
-            reading = await self.client.get_current_reading()
-            status = await self.client.get_current_status()
-        except IOmeterConnectionError as error:
-            raise UpdateFailed(f"Error communicating with IOmeter: {error}") from error
+            async with asyncio.timeout(30):
+                await self._first_data_event.wait()
+        except TimeoutError as err:
+            raise UpdateFailed("Timeout waiting for IOmeter data") from err
+        assert self._reading is not None
+        assert self._status is not None
+        self._update_fw_version(self._status)
+        return IOmeterData(reading=self._reading, status=self._status)
 
+    def _on_new_data(self) -> None:
+        """Called when a new reading or status arrives from SSE."""
+        if self._reading is None or self._status is None:
+            return
+        if not self._first_data_event.is_set():
+            self._first_data_event.set()
+        else:
+            self._update_fw_version(self._status)
+            self.async_set_updated_data(
+                IOmeterData(reading=self._reading, status=self._status)
+            )
+
+    def _on_reading(self, reading: Reading) -> None:
+        """Handle a new reading from the SSE stream."""
+        self._reading = reading
+        self._on_new_data()
+
+    def _on_status(self, status: Status) -> None:
+        """Handle a new status from the SSE stream."""
+        self._status = status
+        self._on_new_data()
+
+    def _on_reading_error(self, err: Exception) -> None:
+        """Log reading stream errors before the library reconnects."""
+        if isinstance(err, IOmeterTimeoutError):
+            _LOGGER.debug("IOmeter reading stream timed out, reconnecting")
+        elif isinstance(err, (IOmeterNoReadingsError, IOmeterConnectionError)):
+            self._async_set_unavailable()
+            _LOGGER.warning("IOmeter reading stream error: %s", err)
+        else:
+            self._async_set_unavailable()
+            _LOGGER.exception("Unexpected error in reading stream")
+
+    def _on_status_error(self, err: Exception) -> None:
+        """Log status stream errors before the library reconnects."""
+        if isinstance(err, IOmeterTimeoutError):
+            _LOGGER.debug("IOmeter status stream timed out, reconnecting")
+        elif isinstance(err, (IOmeterNoStatusError, IOmeterConnectionError)):
+            self._async_set_unavailable()
+            _LOGGER.warning("IOmeter status stream error: %s", err)
+        else:
+            self._async_set_unavailable()
+            _LOGGER.exception("Unexpected error in status stream")
+
+    def _async_set_unavailable(self) -> None:
+        """Mark entities unavailable; skipped before first successful data."""
+        if not self._first_data_event.is_set():
+            return
+        self.last_update_success = False
+        self.async_update_listeners()
+
+    def _update_fw_version(self, status: Status) -> None:
+        """Update device registry if firmware version changed."""
         fw_version = f"{status.device.core.version}/{status.device.bridge.version}"
         if self.current_fw_version and fw_version != self.current_fw_version:
             device_registry = dr.async_get(self.hass)
             device_entry = device_registry.async_get_device_by_identifier(
                 (DOMAIN, status.device.id), self.config_entry.entry_id
             )
-            assert device_entry
-            device_registry.async_update_device(
-                device_entry.id,
-                sw_version=fw_version,
-            )
+            if device_entry:
+                device_registry.async_update_device(
+                    device_entry.id,
+                    sw_version=fw_version,
+                )
         self.current_fw_version = fw_version
-
-        return IOmeterData(reading=reading, status=status)

homeassistant/components/iometer/manifest.json

@@ -5,7 +5,7 @@
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/iometer",
   "integration_type": "device",
-  "iot_class": "local_polling",
+  "iot_class": "local_push",
   "quality_scale": "bronze",
   "requirements": ["iometer==1.0.2"],
   "zeroconf": ["_iometer._tcp.local."]

homeassistant/generated/integrations.json

@@ -3350,7 +3350,7 @@
       "name": "IOmeter",
       "integration_type": "device",
       "config_flow": true,
-      "iot_class": "local_polling"
+      "iot_class": "local_push"
     },
     "ios": {
       "name": "Home Assistant iOS",

tests/components/iometer/__init__.py

@@ -1,6 +1,7 @@
 """Tests for the IOmeter integration."""
 
-from unittest.mock import patch
+from collections.abc import Callable
+from unittest.mock import MagicMock, patch
 
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
@@ -17,3 +18,23 @@ async def setup_platform(
     with patch("homeassistant.components.iometer.PLATFORMS", platforms):
         await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
+
+
+def get_reading_callback(mock: MagicMock) -> Callable:
+    """Get the reading callback registered with the SSE client."""
+    return mock.subscribe_readings.call_args[0][0]
+
+
+def get_status_callback(mock: MagicMock) -> Callable:
+    """Get the status callback registered with the SSE client."""
+    return mock.subscribe_status.call_args[0][0]
+
+
+def get_reading_error_callback(mock: MagicMock) -> Callable:
+    """Get the reading error callback registered with the SSE client."""
+    return mock.subscribe_readings.call_args[0][1]
+
+
+def get_status_error_callback(mock: MagicMock) -> Callable:
+    """Get the status error callback registered with the SSE client."""
+    return mock.subscribe_status.call_args[0][1]

tests/components/iometer/conftest.py

@@ -1,7 +1,7 @@
 """Common fixtures for the IOmeter tests."""
 
 from collections.abc import Generator
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from iometer import Reading, Status
 import pytest
@@ -13,7 +13,7 @@ from tests.common import MockConfigEntry, load_fixture
 
 
 @pytest.fixture
-def mock_setup_entry() -> Generator[AsyncMock]:
+def mock_setup_entry() -> Generator[MagicMock]:
     """Override async_setup_entry."""
     with patch(
         "homeassistant.components.iometer.async_setup_entry",
@@ -23,32 +23,40 @@ def mock_setup_entry() -> Generator[AsyncMock]:
 
 
 @pytest.fixture
-def mock_iometer_client() -> Generator[AsyncMock]:
-    """Mock a new IOmeter client."""
-    with (
-        patch(
-            "homeassistant.components.iometer.IOmeterClient",
-            autospec=True,
-        ) as mock_client,
-        patch(
-            "homeassistant.components.iometer.config_flow.IOmeterClient",
-            new=mock_client,
-        ),
-    ):
-        client = mock_client.return_value
-        client.host = "10.0.0.2"
-        client.get_current_reading.return_value = Reading.from_json(
-            load_fixture("reading.json", DOMAIN)
+def mock_http_client() -> Generator[MagicMock]:
+    """Mock IOmeter HTTP client for config flow."""
+    with patch(
+        "homeassistant.components.iometer.config_flow.IOmeterClient"
+    ) as mock_http_class:
+        http_client = mock_http_class.return_value
+        http_client.get_current_status = AsyncMock(
+            return_value=Status.from_json(load_fixture("status.json", DOMAIN))
         )
-        client.get_current_status.return_value = Status.from_json(
-            load_fixture("status.json", DOMAIN)
-        )
-        yield client
+        yield http_client
+
+
+@pytest.fixture
+def mock_iometer_client(mock_http_client: MagicMock) -> Generator[MagicMock]:
+    """Mock IOmeter SSE client for the coordinator."""
+
+    def subscribe_readings(on_reading, _on_error=None):
+        on_reading(Reading.from_json(load_fixture("reading.json", DOMAIN)))
+        return lambda: None
+
+    def subscribe_status(on_status, _on_error=None):
+        on_status(Status.from_json(load_fixture("status.json", DOMAIN)))
+        return lambda: None
+
+    with patch("homeassistant.components.iometer.IOmeterSSEClient") as mock_sse_class:
+        sse_client = mock_sse_class.return_value
+        sse_client.subscribe_readings.side_effect = subscribe_readings
+        sse_client.subscribe_status.side_effect = subscribe_status
+        yield sse_client
 
 
 @pytest.fixture
 def mock_config_entry() -> MockConfigEntry:
-    """Mock a IOmeter config entry."""
+    """Mock an IOmeter config entry."""
     return MockConfigEntry(
         domain=DOMAIN,
         title="IOmeter-1ISK0000000000",

tests/components/iometer/test_binary_sensor.py

@@ -1,26 +1,27 @@
 """Test the IOmeter binary sensors."""
 
-from datetime import timedelta
-from unittest.mock import AsyncMock
+import json
+from unittest.mock import MagicMock
 
-from freezegun.api import FrozenDateTimeFactory
+from iometer import Status
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
+from homeassistant.components.iometer.const import DOMAIN
 from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNKNOWN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
-from . import setup_platform
+from . import get_status_callback, setup_platform
 
-from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
+from tests.common import MockConfigEntry, async_load_fixture, snapshot_platform
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 async def test_binary_sensors(
     hass: HomeAssistant,
     snapshot: SnapshotAssertion,
-    mock_iometer_client: AsyncMock,
+    mock_iometer_client: MagicMock,
     mock_config_entry: MockConfigEntry,
     entity_registry: er.EntityRegistry,
 ) -> None:
@@ -34,10 +35,9 @@ async def test_binary_sensors(
 async def test_connection_status_sensors(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
-    mock_iometer_client: AsyncMock,
-    freezer: FrozenDateTimeFactory,
+    mock_iometer_client: MagicMock,
 ) -> None:
-    """Test connection status sensor."""
+    """Test connection status sensor updates via SSE."""
     await setup_platform(hass, mock_config_entry, [Platform.BINARY_SENSOR])
 
     assert (
@@ -47,15 +47,9 @@ async def test_connection_status_sensors(
         == STATE_ON
     )
 
-    freezer.tick(delta=timedelta(minutes=1))
-    async_fire_time_changed(hass)
-    await hass.async_block_till_done()
-
-    status = mock_iometer_client.get_current_status.return_value
-    status.device.core.connection_status = "disconnected"
-
-    freezer.tick(delta=timedelta(minutes=1))
-    async_fire_time_changed(hass)
+    status_data = json.loads(await async_load_fixture(hass, "status.json", DOMAIN))
+    status_data["device"]["core"]["connectionStatus"] = "disconnected"
+    get_status_callback(mock_iometer_client)(Status.from_json(json.dumps(status_data)))
     await hass.async_block_till_done()
 
     assert (
@@ -70,10 +64,9 @@ async def test_connection_status_sensors(
 async def test_attachment_status_sensors(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
-    mock_iometer_client: AsyncMock,
-    freezer: FrozenDateTimeFactory,
+    mock_iometer_client: MagicMock,
 ) -> None:
-    """Test connection status sensor."""
+    """Test attachment status sensor updates via SSE."""
     await setup_platform(hass, mock_config_entry, [Platform.BINARY_SENSOR])
 
     assert (
@@ -83,15 +76,9 @@ async def test_attachment_status_sensors(
         == STATE_ON
     )
 
-    freezer.tick(delta=timedelta(minutes=1))
-    async_fire_time_changed(hass)
-    await hass.async_block_till_done()
-
-    status = mock_iometer_client.get_current_status.return_value
-    status.device.core.attachment_status = "detached"
-
-    freezer.tick(delta=timedelta(minutes=1))
-    async_fire_time_changed(hass)
+    status_data = json.loads(await async_load_fixture(hass, "status.json", DOMAIN))
+    status_data["device"]["core"]["attachmentStatus"] = "detached"
+    get_status_callback(mock_iometer_client)(Status.from_json(json.dumps(status_data)))
     await hass.async_block_till_done()
 
     assert (
@@ -103,13 +90,12 @@ async def test_attachment_status_sensors(
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
-async def test_attachment_status_sensors_unkown(
+async def test_attachment_status_sensors_unknown(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
-    mock_iometer_client: AsyncMock,
-    freezer: FrozenDateTimeFactory,
+    mock_iometer_client: MagicMock,
 ) -> None:
-    """Test connection status sensor."""
+    """Test attachment status sensor shows unknown state via SSE."""
     await setup_platform(hass, mock_config_entry, [Platform.BINARY_SENSOR])
 
     assert (
@@ -119,15 +105,9 @@ async def test_attachment_status_sensors_unkown(
         == STATE_ON
     )
 
-    freezer.tick(delta=timedelta(minutes=1))
-    async_fire_time_changed(hass)
-    await hass.async_block_till_done()
-
-    status = mock_iometer_client.get_current_status.return_value
-    status.device.core.attachment_status = None
-
-    freezer.tick(delta=timedelta(minutes=1))
-    async_fire_time_changed(hass)
+    status_data = json.loads(await async_load_fixture(hass, "status.json", DOMAIN))
+    del status_data["device"]["core"]["attachmentStatus"]
+    get_status_callback(mock_iometer_client)(Status.from_json(json.dumps(status_data)))
     await hass.async_block_till_done()
 
     assert (

tests/components/iometer/test_config_flow.py

@@ -1,9 +1,14 @@
 """Test the IOmeter config flow."""
 
 from ipaddress import ip_address
-from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
 
-from iometer import IOmeterConnectionError, IOmeterNoReadingsError, IOmeterNoStatusError
+from iometer import (
+    IOmeterConnectionError,
+    IOmeterNoStatusError,
+    IOmeterTimeoutError,
+    Status,
+)
 import pytest
 
 from homeassistant.components.iometer.const import DOMAIN
@@ -13,7 +18,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 from homeassistant.helpers.service_info.zeroconf import ZeroconfServiceInfo
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_load_fixture
 
 IP_ADDRESS = "10.0.0.2"
 IOMETER_DEVICE_ID = "658c2b34-2017-45f2-a12b-731235f8bb97"
@@ -29,16 +34,16 @@ ZEROCONF_DISCOVERY = ZeroconfServiceInfo(
 )
 
 
+@pytest.mark.usefixtures("mock_setup_entry")
 async def test_user_flow(
     hass: HomeAssistant,
-    mock_iometer_client: AsyncMock,
+    mock_http_client: MagicMock,
 ) -> None:
     """Test full user configuration flow."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": SOURCE_USER},
     )
-    await hass.async_block_till_done()
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "user"
 
@@ -46,7 +51,6 @@ async def test_user_flow(
         result["flow_id"],
         user_input={CONF_HOST: IP_ADDRESS},
     )
-
     await hass.async_block_till_done()
     assert result["type"] is FlowResultType.CREATE_ENTRY
     assert result["title"] == "IOmeter 1ISK0000000000"
@@ -54,9 +58,10 @@ async def test_user_flow(
     assert result["result"].unique_id == IOMETER_DEVICE_ID
 
 
+@pytest.mark.usefixtures("mock_setup_entry")
 async def test_zeroconf_flow(
     hass: HomeAssistant,
-    mock_iometer_client: AsyncMock,
+    mock_http_client: MagicMock,
 ) -> None:
     """Test zeroconf flow."""
     result = await hass.config_entries.flow.async_init(
@@ -64,7 +69,6 @@ async def test_zeroconf_flow(
         context={"source": SOURCE_ZEROCONF},
         data=ZEROCONF_DISCOVERY,
     )
-    await hass.async_block_till_done()
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "zeroconf_confirm"
 
@@ -95,23 +99,22 @@ async def test_zeroconf_flow_abort_duplicate(
 
 
 @pytest.mark.parametrize(
-    ("method_name", "exception", "reason"),
+    ("exception", "reason"),
     [
-        ("get_current_status", IOmeterConnectionError(), "cannot_connect"),
-        ("get_current_status", IOmeterNoStatusError(), "no_status"),
-        ("get_current_reading", IOmeterNoReadingsError(), "no_readings"),
+        (IOmeterConnectionError(), "cannot_connect"),
+        (IOmeterTimeoutError(), "cannot_connect"),
+        (IOmeterNoStatusError(), "no_status"),
     ],
-    ids=["status-connection", "status-missing", "reading-missing"],
+    ids=["connection-error", "timeout", "status-missing"],
 )
 async def test_zeroconf_flow_abort_errors(
     hass: HomeAssistant,
-    mock_iometer_client: AsyncMock,
-    method_name: str,
+    mock_http_client: MagicMock,
     exception: Exception,
     reason: str,
 ) -> None:
-    """Test zeroconf flow aborts when the client raises an exception."""
-    getattr(mock_iometer_client, method_name).side_effect = exception
+    """Test zeroconf flow aborts when the HTTP client raises an exception."""
+    mock_http_client.get_current_status.side_effect = exception
 
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
@@ -123,31 +126,51 @@ async def test_zeroconf_flow_abort_errors(
     assert result["reason"] == reason
 
 
+async def test_zeroconf_flow_abort_no_meter(
+    hass: HomeAssistant,
+    mock_http_client: MagicMock,
+) -> None:
+    """Test zeroconf flow aborts when the status contains no meter info."""
+    mock_status = MagicMock()
+    mock_status.meter = None
+    mock_http_client.get_current_status.return_value = mock_status
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": SOURCE_ZEROCONF},
+        data=ZEROCONF_DISCOVERY,
+    )
+    await hass.async_block_till_done()
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "no_readings"
+
+
 @pytest.mark.parametrize(
-    ("method_name", "exception", "error_key"),
+    ("exception", "error_key"),
     [
-        ("get_current_status", IOmeterConnectionError(), "cannot_connect"),
-        ("get_current_status", IOmeterNoStatusError(), "no_status"),
-        ("get_current_reading", IOmeterNoReadingsError(), "no_readings"),
+        (IOmeterConnectionError(), "cannot_connect"),
+        (IOmeterTimeoutError(), "cannot_connect"),
+        (IOmeterNoStatusError(), "no_status"),
     ],
-    ids=["status-connection", "status-missing", "reading-missing"],
+    ids=["connection-error", "timeout", "status-missing"],
 )
 @pytest.mark.usefixtures("mock_setup_entry")
 async def test_user_flow_errors(
     hass: HomeAssistant,
-    mock_iometer_client: AsyncMock,
-    method_name: str,
+    mock_http_client: MagicMock,
     exception: Exception,
     error_key: str,
 ) -> None:
-    """Test user flow returns errors for client exceptions."""
-    getattr(mock_iometer_client, method_name).side_effect = exception
+    """Test user flow shows errors for HTTP client exceptions and recovers on retry."""
+    valid_status = Status.from_json(
+        await async_load_fixture(hass, "status.json", DOMAIN)
+    )
+    mock_http_client.get_current_status.side_effect = [exception, valid_status]
 
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": SOURCE_USER},
     )
-    await hass.async_block_till_done()
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "user"
 
@@ -159,7 +182,41 @@ async def test_user_flow_errors(
     assert result["type"] is FlowResultType.FORM
     assert result["errors"] == {"base": error_key}
 
-    getattr(mock_iometer_client, method_name).side_effect = None
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_HOST: IP_ADDRESS},
+    )
+    await hass.async_block_till_done()
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+
+@pytest.mark.usefixtures("mock_setup_entry")
+async def test_user_flow_no_meter_error(
+    hass: HomeAssistant,
+    mock_http_client: MagicMock,
+) -> None:
+    """Test user flow shows error when status contains no meter info."""
+    mock_status = MagicMock()
+    mock_status.meter = None
+    valid_status = Status.from_json(
+        await async_load_fixture(hass, "status.json", DOMAIN)
+    )
+    mock_http_client.get_current_status.side_effect = [mock_status, valid_status]
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": SOURCE_USER},
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "user"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_HOST: IP_ADDRESS},
+    )
+    await hass.async_block_till_done()
+    assert result["type"] is FlowResultType.FORM
+    assert result["errors"] == {"base": "no_readings"}
 
     result = await hass.config_entries.flow.async_configure(
         result["flow_id"],
@@ -172,7 +229,7 @@ async def test_user_flow_errors(
 @pytest.mark.usefixtures("mock_setup_entry")
 async def test_flow_abort_duplicate(
     hass: HomeAssistant,
-    mock_iometer_client: AsyncMock,
+    mock_http_client: MagicMock,
     mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test duplicate flow."""
@@ -182,7 +239,6 @@ async def test_flow_abort_duplicate(
         DOMAIN,
         context={"source": SOURCE_USER},
     )
-    await hass.async_block_till_done()
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "user"
 
@@ -191,6 +247,5 @@ async def test_flow_abort_duplicate(
         {CONF_HOST: IP_ADDRESS},
     )
     await hass.async_block_till_done()
-
     assert result["type"] is FlowResultType.ABORT
     assert result["reason"] == "already_configured"

tests/components/iometer/test_init.py

@@ -1,10 +1,19 @@
 """Tests for the IOmeter integration."""
 
-from datetime import timedelta
-from unittest.mock import AsyncMock
+import asyncio
+import contextlib
+import json
+import logging
+from unittest.mock import AsyncMock, MagicMock, patch
 
-from freezegun.api import FrozenDateTimeFactory
-from iometer import IOmeterConnectionError
+from iometer import (
+    IOmeterConnectionError,
+    IOmeterNoReadingsError,
+    IOmeterNoStatusError,
+    IOmeterTimeoutError,
+    Status,
+)
+import pytest
 
 from homeassistant.components.iometer.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
@@ -12,19 +21,23 @@ from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 
-from . import setup_platform
+from . import (
+    get_reading_error_callback,
+    get_status_callback,
+    get_status_error_callback,
+    setup_platform,
+)
 
-from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.common import MockConfigEntry, async_load_fixture
 
 
 async def test_new_firmware_version(
     hass: HomeAssistant,
-    mock_iometer_client: AsyncMock,
+    mock_iometer_client: MagicMock,
     mock_config_entry: MockConfigEntry,
     device_registry: dr.DeviceRegistry,
-    freezer: FrozenDateTimeFactory,
 ) -> None:
-    """Test device registry integration."""
+    """Test device registry is updated when firmware version changes via SSE."""
     assert mock_config_entry.unique_id is not None
 
     await setup_platform(hass, mock_config_entry, [Platform.SENSOR])
@@ -33,13 +46,13 @@ async def test_new_firmware_version(
     )
     assert device_entry is not None
     assert device_entry.sw_version == "build-58/build-65"
-    mock_iometer_client.get_current_status.return_value.device.core.version = "build-62"
-    mock_iometer_client.get_current_status.return_value.device.bridge.version = (
-        "build-69"
-    )
-    freezer.tick(timedelta(minutes=1))
-    async_fire_time_changed(hass)
+
+    status_data = json.loads(await async_load_fixture(hass, "status.json", DOMAIN))
+    status_data["device"]["core"]["version"] = "build-62"
+    status_data["device"]["bridge"]["version"] = "build-69"
+    get_status_callback(mock_iometer_client)(Status.from_json(json.dumps(status_data)))
     await hass.async_block_till_done()
+
     device_entry = device_registry.async_get_device(
         identifiers={(DOMAIN, mock_config_entry.unique_id)}
     )
@@ -47,19 +60,160 @@ async def test_new_firmware_version(
     assert device_entry.sw_version == "build-62/build-69"
 
 
-async def test_async_setup_entry_connection_error(
+async def test_first_data_timeout(
     hass: HomeAssistant,
-    mock_iometer_client: AsyncMock,
+    mock_iometer_client: MagicMock,
     mock_config_entry: MockConfigEntry,
 ) -> None:
-    """Test async_setup_entry raises ConfigEntryNotReady on connection error."""
+    """Test setup retries when the 30s timeout waiting for first SSE data expires."""
+    mock_config_entry.add_to_hass(hass)
+
+    mock_timeout = MagicMock()
+    mock_timeout.return_value.__aenter__ = AsyncMock(side_effect=TimeoutError)
+    mock_timeout.return_value.__aexit__ = AsyncMock(return_value=False)
+
+    with patch(
+        "homeassistant.components.iometer.coordinator.asyncio.timeout",
+        mock_timeout,
+    ):
+        result = await hass.config_entries.async_setup(mock_config_entry.entry_id)
+
+    assert not result
+    assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+@pytest.mark.parametrize(
+    ("exception", "expected_log"),
+    [
+        pytest.param(IOmeterTimeoutError("t"), "timed out", id="timeout"),
+        pytest.param(IOmeterNoReadingsError("n"), "stream error", id="no-readings"),
+        pytest.param(
+            IOmeterConnectionError("c"), "stream error", id="connection-error"
+        ),
+        pytest.param(
+            RuntimeError("u"), "Unexpected error in reading stream", id="unexpected"
+        ),
+    ],
+)
+async def test_reading_error_callback(
+    hass: HomeAssistant,
+    mock_iometer_client: MagicMock,
+    mock_config_entry: MockConfigEntry,
+    exception: Exception,
+    expected_log: str,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test reading error callback logs correctly before the library reconnects."""
+    await setup_platform(hass, mock_config_entry, [Platform.SENSOR])
+
+    with caplog.at_level(logging.DEBUG, logger="homeassistant.components.iometer"):
+        get_reading_error_callback(mock_iometer_client)(exception)
+
+    assert expected_log in caplog.text
+
+
+@pytest.mark.parametrize(
+    ("exception", "expected_log"),
+    [
+        pytest.param(IOmeterTimeoutError("t"), "timed out", id="timeout"),
+        pytest.param(IOmeterNoStatusError("n"), "stream error", id="no-status"),
+        pytest.param(
+            IOmeterConnectionError("c"), "stream error", id="connection-error"
+        ),
+        pytest.param(
+            RuntimeError("u"), "Unexpected error in status stream", id="unexpected"
+        ),
+    ],
+)
+async def test_status_error_callback(
+    hass: HomeAssistant,
+    mock_iometer_client: MagicMock,
+    mock_config_entry: MockConfigEntry,
+    exception: Exception,
+    expected_log: str,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test status error callback logs correctly before the library reconnects."""
+    await setup_platform(hass, mock_config_entry, [Platform.SENSOR])
+
+    with caplog.at_level(logging.DEBUG, logger="homeassistant.components.iometer"):
+        get_status_error_callback(mock_iometer_client)(exception)
+
+    assert expected_log in caplog.text
+
+
+async def test_error_before_first_data_does_not_mark_unavailable(
+    hass: HomeAssistant,
+    mock_iometer_client: MagicMock,
+    mock_config_entry: MockConfigEntry,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that stream errors before first data do not mark entities unavailable."""
+    subscribed = asyncio.Event()
+
+    def subscribe_readings_noop(*_):
+        subscribed.set()
+        return lambda: None
+
+    mock_iometer_client.subscribe_readings.side_effect = subscribe_readings_noop
+    mock_iometer_client.subscribe_status.side_effect = lambda *_: lambda: None
 
     mock_config_entry.add_to_hass(hass)
-    mock_iometer_client.get_current_status.side_effect = IOmeterConnectionError(
-        "cannot connect"
-    )
 
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    with patch("homeassistant.components.iometer.PLATFORMS", []):
+        setup_task = asyncio.get_running_loop().create_task(
+            hass.config_entries.async_setup(mock_config_entry.entry_id)
+        )
+        await asyncio.wait_for(subscribed.wait(), timeout=5.0)
 
-    assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
-    assert mock_iometer_client.get_current_status.await_count == 1
+    with caplog.at_level(logging.WARNING, logger="homeassistant.components.iometer"):
+        get_reading_error_callback(mock_iometer_client)(IOmeterConnectionError("err"))
+
+    assert "stream error" in caplog.text
+    assert "Update failed" not in caplog.text
+
+    setup_task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await setup_task
+
+
+@pytest.mark.usefixtures("mock_iometer_client")
+async def test_async_unload_entry(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test that unloading an entry succeeds and cleans up the coordinator."""
+    await setup_platform(hass, mock_config_entry, [Platform.SENSOR])
+    assert await hass.config_entries.async_unload(mock_config_entry.entry_id)
+    assert mock_config_entry.state is ConfigEntryState.NOT_LOADED
+
+
+async def test_async_stop_calls_cancel_on_unload(
+    hass: HomeAssistant,
+    mock_iometer_client: MagicMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test that unloading calls the cancel functions returned by the library."""
+    readings_cancel = MagicMock()
+    status_cancel = MagicMock()
+
+    original_readings = mock_iometer_client.subscribe_readings.side_effect
+    original_status = mock_iometer_client.subscribe_status.side_effect
+
+    def subscribe_readings_with_cancel(*args):
+        original_readings(*args)
+        return readings_cancel
+
+    def subscribe_status_with_cancel(*args):
+        original_status(*args)
+        return status_cancel
+
+    mock_iometer_client.subscribe_readings.side_effect = subscribe_readings_with_cancel
+    mock_iometer_client.subscribe_status.side_effect = subscribe_status_with_cancel
+
+    await setup_platform(hass, mock_config_entry, [Platform.SENSOR])
+    assert await hass.config_entries.async_unload(mock_config_entry.entry_id)
+    assert mock_config_entry.state is ConfigEntryState.NOT_LOADED
+
+    readings_cancel.assert_called_once()
+    status_cancel.assert_called_once()

tests/components/iometer/test_sensor.py

@@ -1,5 +1,7 @@
 """Test the sensors provided by the Powerfox integration."""
 
+from unittest.mock import MagicMock
+
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
@@ -10,13 +12,12 @@ from homeassistant.helpers import entity_registry as er
 from . import setup_platform
 
 from tests.common import MockConfigEntry, snapshot_platform
-from tests.components.conftest import AsyncMock
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 async def test_all_sensors(
     hass: HomeAssistant,
-    mock_iometer_client: AsyncMock,
+    mock_iometer_client: MagicMock,
     mock_config_entry: MockConfigEntry,
     entity_registry: er.EntityRegistry,
     snapshot: SnapshotAssertion,