Migrate system health to use integration_platform (#42785)

homeassistant/components/lovelace/__init__.py

@@ -140,8 +140,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType):
         websocket.websocket_lovelace_dashboards
     )
 
-    hass.components.system_health.async_register_info(DOMAIN, system_health_info)
-
     hass.data[DOMAIN] = {
         # We store a dictionary mapping url_path: config. None is the default.
         "dashboards": {None: default_config},
@@ -233,14 +231,6 @@ async def create_yaml_resource_col(hass, yaml_resources):
     return resources.ResourceYAMLCollection(yaml_resources or [])
 
 
-async def system_health_info(hass):
-    """Get info for the info page."""
-    health_info = {"dashboards": len(hass.data[DOMAIN]["dashboards"])}
-    health_info.update(await hass.data[DOMAIN]["dashboards"][None].async_get_info())
-    health_info.update(await hass.data[DOMAIN]["resources"].async_get_info())
-    return health_info
-
-
 @callback
 def _register_panel(hass, url_path, mode, config, update):
     """Register a panel."""

homeassistant/components/lovelace/system_health.py

@@ -0,0 +1,21 @@
+"""Provide info to system health."""
+from homeassistant.components import system_health
+from homeassistant.core import HomeAssistant, callback
+
+from .const import DOMAIN
+
+
+@callback
+def async_register(
+    hass: HomeAssistant, register: system_health.RegisterSystemHealth
+) -> None:
+    """Register system health callbacks."""
+    register.async_register_info(system_health_info)
+
+
+async def system_health_info(hass):
+    """Get info for the info page."""
+    health_info = {"dashboards": len(hass.data[DOMAIN]["dashboards"])}
+    health_info.update(await hass.data[DOMAIN]["dashboards"][None].async_get_info())
+    health_info.update(await hass.data[DOMAIN]["resources"].async_get_info())
+    return health_info

homeassistant/components/system_health/__init__.py

@@ -1,6 +1,6 @@
 """Support for System health ."""
 import asyncio
-from collections import OrderedDict
+import dataclasses
 import logging
 from typing import Callable, Dict
 
@@ -8,8 +8,9 @@ import async_timeout
 import voluptuous as vol
 
 from homeassistant.components import websocket_api
-from homeassistant.core import callback
-from homeassistant.helpers.typing import ConfigType, HomeAssistantType
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import integration_platform
+from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import bind_hass
 
 _LOGGER = logging.getLogger(__name__)
@@ -22,21 +23,38 @@ INFO_CALLBACK_TIMEOUT = 5
 @bind_hass
 @callback
 def async_register_info(
-    hass: HomeAssistantType,
+    hass: HomeAssistant,
     domain: str,
-    info_callback: Callable[[HomeAssistantType], Dict],
+    info_callback: Callable[[HomeAssistant], Dict],
 ):
-    """Register an info callback."""
-    data = hass.data.setdefault(DOMAIN, OrderedDict()).setdefault("info", OrderedDict())
-    data[domain] = info_callback
+    """Register an info callback.
+
+    Deprecated.
+    """
+    _LOGGER.warning(
+        "system_health.async_register_info is deprecated. Add a system_health platform instead."
+    )
+    hass.data.setdefault(DOMAIN, {}).setdefault("info", {})
+    RegisterSystemHealth(hass, domain).async_register_info(info_callback)
 
 
-async def async_setup(hass: HomeAssistantType, config: ConfigType):
+async def async_setup(hass: HomeAssistant, config: ConfigType):
     """Set up the System Health component."""
     hass.components.websocket_api.async_register_command(handle_info)
+    hass.data.setdefault(DOMAIN, {"info": {}})
+
+    await integration_platform.async_process_integration_platforms(
+        hass, DOMAIN, _register_system_health_platform
+    )
+
     return True
 
 
+async def _register_system_health_platform(hass, integration_domain, platform):
+    """Register a system health platform."""
+    platform.async_register(hass, RegisterSystemHealth(hass, integration_domain))
+
+
 async def _info_wrapper(hass, info_callback):
     """Wrap info callback."""
     try:
@@ -52,11 +70,11 @@ async def _info_wrapper(hass, info_callback):
 @websocket_api.async_response
 @websocket_api.websocket_command({vol.Required("type"): "system_health/info"})
 async def handle_info(
-    hass: HomeAssistantType, connection: websocket_api.ActiveConnection, msg: Dict
+    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: Dict
 ):
     """Handle an info request."""
     info_callbacks = hass.data.get(DOMAIN, {}).get("info", {})
-    data = OrderedDict()
+    data = {}
     data["homeassistant"] = await hass.helpers.system_info.async_get_system_info()
 
     if info_callbacks:
@@ -72,3 +90,19 @@ async def handle_info(
             data[domain] = domain_data
 
     connection.send_message(websocket_api.result_message(msg["id"], data))
+
+
+@dataclasses.dataclass(frozen=True)
+class RegisterSystemHealth:
+    """Helper class to allow platforms to register."""
+
+    hass: HomeAssistant
+    domain: str
+
+    @callback
+    def async_register_info(
+        self,
+        info_callback: Callable[[HomeAssistant], Dict],
+    ):
+        """Register an info callback."""
+        self.hass.data[DOMAIN]["info"][self.domain] = info_callback

tests/components/lovelace/test_dashboard.py

@@ -6,11 +6,7 @@ from homeassistant.components.lovelace import const, dashboard
 from homeassistant.setup import async_setup_component
 
 from tests.async_mock import patch
-from tests.common import (
-    assert_setup_component,
-    async_capture_events,
-    get_system_health_info,
-)
+from tests.common import assert_setup_component, async_capture_events
 
 
 async def test_lovelace_from_storage(hass, hass_ws_client, hass_storage):
@@ -160,48 +156,6 @@ async def test_lovelace_from_yaml(hass, hass_ws_client):
     assert len(events) == 1
 
 
-async def test_system_health_info_autogen(hass):
-    """Test system health info endpoint."""
-    assert await async_setup_component(hass, "lovelace", {})
-    info = await get_system_health_info(hass, "lovelace")
-    assert info == {"dashboards": 1, "mode": "auto-gen", "resources": 0}
-
-
-async def test_system_health_info_storage(hass, hass_storage):
-    """Test system health info endpoint."""
-    hass_storage[dashboard.CONFIG_STORAGE_KEY_DEFAULT] = {
-        "key": "lovelace",
-        "version": 1,
-        "data": {"config": {"resources": [], "views": []}},
-    }
-    assert await async_setup_component(hass, "lovelace", {})
-    info = await get_system_health_info(hass, "lovelace")
-    assert info == {"dashboards": 1, "mode": "storage", "resources": 0, "views": 0}
-
-
-async def test_system_health_info_yaml(hass):
-    """Test system health info endpoint."""
-    assert await async_setup_component(hass, "lovelace", {"lovelace": {"mode": "YAML"}})
-    with patch(
-        "homeassistant.components.lovelace.dashboard.load_yaml",
-        return_value={"views": [{"cards": []}]},
-    ):
-        info = await get_system_health_info(hass, "lovelace")
-    assert info == {"dashboards": 1, "mode": "yaml", "resources": 0, "views": 1}
-
-
-async def test_system_health_info_yaml_not_found(hass):
-    """Test system health info endpoint."""
-    assert await async_setup_component(hass, "lovelace", {"lovelace": {"mode": "YAML"}})
-    info = await get_system_health_info(hass, "lovelace")
-    assert info == {
-        "dashboards": 1,
-        "mode": "yaml",
-        "error": "{} not found".format(hass.config.path("ui-lovelace.yaml")),
-        "resources": 0,
-    }
-
-
 @pytest.mark.parametrize("url_path", ("test-panel", "test-panel-no-sidebar"))
 async def test_dashboard_from_yaml(hass, hass_ws_client, url_path):
     """Test we load lovelace dashboard config from yaml."""

tests/components/lovelace/test_system_health.py

@@ -0,0 +1,52 @@
+"""Tests for Lovelace system health."""
+from homeassistant.components.lovelace import dashboard
+from homeassistant.setup import async_setup_component
+
+from tests.async_mock import patch
+from tests.common import get_system_health_info
+
+
+async def test_system_health_info_autogen(hass):
+    """Test system health info endpoint."""
+    assert await async_setup_component(hass, "lovelace", {})
+    assert await async_setup_component(hass, "system_health", {})
+    info = await get_system_health_info(hass, "lovelace")
+    assert info == {"dashboards": 1, "mode": "auto-gen", "resources": 0}
+
+
+async def test_system_health_info_storage(hass, hass_storage):
+    """Test system health info endpoint."""
+    assert await async_setup_component(hass, "system_health", {})
+    hass_storage[dashboard.CONFIG_STORAGE_KEY_DEFAULT] = {
+        "key": "lovelace",
+        "version": 1,
+        "data": {"config": {"resources": [], "views": []}},
+    }
+    assert await async_setup_component(hass, "lovelace", {})
+    info = await get_system_health_info(hass, "lovelace")
+    assert info == {"dashboards": 1, "mode": "storage", "resources": 0, "views": 0}
+
+
+async def test_system_health_info_yaml(hass):
+    """Test system health info endpoint."""
+    assert await async_setup_component(hass, "system_health", {})
+    assert await async_setup_component(hass, "lovelace", {"lovelace": {"mode": "YAML"}})
+    with patch(
+        "homeassistant.components.lovelace.dashboard.load_yaml",
+        return_value={"views": [{"cards": []}]},
+    ):
+        info = await get_system_health_info(hass, "lovelace")
+    assert info == {"dashboards": 1, "mode": "yaml", "resources": 0, "views": 1}
+
+
+async def test_system_health_info_yaml_not_found(hass):
+    """Test system health info endpoint."""
+    assert await async_setup_component(hass, "system_health", {})
+    assert await async_setup_component(hass, "lovelace", {"lovelace": {"mode": "YAML"}})
+    info = await get_system_health_info(hass, "lovelace")
+    assert info == {
+        "dashboards": 1,
+        "mode": "yaml",
+        "error": "{} not found".format(hass.config.path("ui-lovelace.yaml")),
+        "resources": 0,
+    }

tests/components/system_health/test_init.py

@@ -1,19 +1,19 @@
 """Tests for the system health component init."""
 import asyncio
-from unittest.mock import Mock
 
 import pytest
 
 from homeassistant.setup import async_setup_component
 
-from tests.common import mock_coro
+from tests.async_mock import AsyncMock, Mock
+from tests.common import get_system_health_info, mock_platform
 
 
 @pytest.fixture
 def mock_system_info(hass):
     """Mock system info."""
-    hass.helpers.system_info.async_get_system_info = Mock(
-        return_value=mock_coro({"hello": True})
+    hass.helpers.system_info.async_get_system_info = AsyncMock(
+        return_value={"hello": True}
     )
 
 
@@ -51,6 +51,9 @@ async def test_info_endpoint_register_callback(hass, hass_ws_client, mock_system
     data = data["lovelace"]
     assert data == {"storage": "YAML"}
 
+    # Test our test helper works
+    assert await get_system_health_info(hass, "lovelace") == {"storage": "YAML"}
+
 
 async def test_info_endpoint_register_callback_timeout(
     hass, hass_ws_client, mock_system_info
@@ -94,3 +97,20 @@ async def test_info_endpoint_register_callback_exc(
     assert len(data) == 2
     data = data["lovelace"]
     assert data == {"error": "TEST ERROR"}
+
+
+async def test_platform_loading(hass):
+    """Test registering via platform."""
+    hass.config.components.add("fake_integration")
+    mock_platform(
+        hass,
+        "fake_integration.system_health",
+        Mock(
+            async_register=lambda hass, register: register.async_register_info(
+                AsyncMock(return_value={"hello": "info"})
+            )
+        ),
+    )
+
+    assert await async_setup_component(hass, "system_health", {})
+    assert await get_system_health_info(hass, "fake_integration") == {"hello": "info"}