Add system health streaming updates

homeassistant/components/cloud/system_health.py

@@ -0,0 +1,48 @@
+"""Provide info to system health."""
+from hass_nabucasa import Cloud
+from yarl import URL
+
+from homeassistant.components import system_health
+from homeassistant.core import HomeAssistant, callback
+
+from .client import CloudClient
+from .const import DOMAIN
+
+
+@callback
+def async_register(
+    hass: HomeAssistant, register: system_health.SystemHealthRegistration
+) -> None:
+    """Register system health callbacks."""
+    register.async_register_info(system_health_info, "/config/cloud")
+
+
+async def system_health_info(hass):
+    """Get info for the info page."""
+    cloud: Cloud = hass.data[DOMAIN]
+    client: CloudClient = cloud.client
+
+    data = {
+        "logged_in": cloud.is_logged_in,
+    }
+
+    if cloud.is_logged_in:
+        data["subscription_expiration"] = cloud.expiration_date
+        data["relayer_connected"] = cloud.is_connected
+        data["remote_enabled"] = client.prefs.remote_enabled
+        data["remote_connected"] = cloud.remote.is_connected
+        data["alexa_enabled"] = client.prefs.alexa_enabled
+        data["google_enabled"] = client.prefs.google_enabled
+
+    data["can_reach_cert_server"] = system_health.async_check_can_reach_url(
+        hass, cloud.acme_directory_server
+    )
+    data["can_reach_cloud_auth"] = system_health.async_check_can_reach_url(
+        hass,
+        f"https://cognito-idp.{cloud.region}.amazonaws.com/{cloud.user_pool_id}/.well-known/jwks.json",
+    )
+    data["can_reach_cloud"] = system_health.async_check_can_reach_url(
+        hass, URL(cloud.relayer).with_scheme("https").with_path("/status")
+    )
+
+    return data

homeassistant/components/lovelace/system_health.py

@@ -7,10 +7,10 @@
 
 @callback
 def async_register(
-    hass: HomeAssistant, register: system_health.RegisterSystemHealth
+    hass: HomeAssistant, register: system_health.SystemHealthRegistration
 ) -> None:
     """Register system health callbacks."""
-    register.async_register_info(system_health_info)
+    register.async_register_info(system_health_info, "/config/lovelace")
 
 
 async def system_health_info(hass):

homeassistant/components/system_health/__init__.py

@@ -1,15 +1,17 @@
 """Support for System health ."""
 import asyncio
 import dataclasses
+from datetime import datetime
 import logging
-from typing import Callable, Dict
+from typing import Awaitable, Callable, Dict, Optional
 
+import aiohttp
 import async_timeout
 import voluptuous as vol
 
 from homeassistant.components import websocket_api
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import integration_platform
+from homeassistant.helpers import aiohttp_client, integration_platform
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import bind_hass
 
@@ -34,14 +36,14 @@ def async_register_info(
     _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)
+    hass.data.setdefault(DOMAIN, {})
+    SystemHealthRegistration(hass, domain).async_register_info(info_callback)
 
 
 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": {}})
+    hass.data.setdefault(DOMAIN, {})
 
     await integration_platform.async_process_integration_platforms(
         hass, DOMAIN, _register_system_health_platform
@@ -52,57 +54,169 @@ async def async_setup(hass: HomeAssistant, config: ConfigType):
 
 async def _register_system_health_platform(hass, integration_domain, platform):
     """Register a system health platform."""
-    platform.async_register(hass, RegisterSystemHealth(hass, integration_domain))
+    platform.async_register(hass, SystemHealthRegistration(hass, integration_domain))
 
 
-async def _info_wrapper(hass, info_callback):
-    """Wrap info callback."""
+async def get_integration_info(
+    hass: HomeAssistant, registration: "SystemHealthRegistration"
+):
+    """Get integration system health."""
     try:
         with async_timeout.timeout(INFO_CALLBACK_TIMEOUT):
-            return await info_callback(hass)
+            data = await registration.info_callback(hass)
     except asyncio.TimeoutError:
-        return {"error": "Fetching info timed out"}
-    except Exception as err:  # pylint: disable=broad-except
+        data = {"error": {"type": "failed", "error": "timeout"}}
+    except Exception:  # pylint: disable=broad-except
         _LOGGER.exception("Error fetching info")
-        return {"error": str(err)}
+        data = {"error": {"type": "failed", "error": "unknown"}}
+
+    result = {"info": data}
+
+    if registration.manage_url:
+        result["manage_url"] = registration.manage_url
+
+    return result
+
+
+@callback
+def _format_value(val):
+    """Format a system health value."""
+    if isinstance(val, datetime):
+        return {"value": val.isoformat(), "type": "date"}
+    return val
 
 
 @websocket_api.async_response
 @websocket_api.websocket_command({vol.Required("type"): "system_health/info"})
 async def handle_info(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: Dict
 ):
-    """Handle an info request."""
-    info_callbacks = hass.data.get(DOMAIN, {}).get("info", {})
+    """Handle an info request via a subscription."""
+    registrations: Dict[str, SystemHealthRegistration] = hass.data[DOMAIN]
     data = {}
-    data["homeassistant"] = await hass.helpers.system_info.async_get_system_info()
-
-    if info_callbacks:
-        for domain, domain_data in zip(
-            info_callbacks,
-            await asyncio.gather(
-                *(
-                    _info_wrapper(hass, info_callback)
-                    for info_callback in info_callbacks.values()
-                )
-            ),
-        ):
-            data[domain] = domain_data
+    data["homeassistant"] = {
+        "info": await hass.helpers.system_info.async_get_system_info()
+    }
+
+    pending_info = {}
+
+    for domain, domain_data in zip(
+        registrations,
+        await asyncio.gather(
+            *(
+                get_integration_info(hass, registration)
+                for registration in registrations.values()
+            )
+        ),
+    ):
+        for key, value in domain_data["info"].items():
+            if asyncio.iscoroutine(value):
+                value = asyncio.create_task(value)
+            if isinstance(value, asyncio.Task):
+                pending_info[(domain, key)] = value
+                domain_data["info"][key] = {"type": "pending"}
+            else:
+                domain_data["info"][key] = _format_value(value)
+
+        data[domain] = domain_data
+
+    # Confirm subscription
+    connection.send_result(msg["id"])
+
+    stop_event = asyncio.Event()
+    connection.subscriptions[msg["id"]] = stop_event.set
+
+    # Send initial data
+    connection.send_message(
+        websocket_api.messages.event_message(
+            msg["id"], {"type": "initial", "data": data}
+        )
+    )
 
-    connection.send_message(websocket_api.result_message(msg["id"], data))
+    # If nothing pending, wrap it up.
+    if not pending_info:
+        connection.send_message(
+            websocket_api.messages.event_message(msg["id"], {"type": "finish"})
+        )
+        return
+
+    tasks = [asyncio.create_task(stop_event.wait()), *pending_info.values()]
+    pending_lookup = {val: key for key, val in pending_info.items()}
+
+    # One task is the stop_event.wait() and is always there
+    while len(tasks) > 1 and not stop_event.is_set():
+        # Wait for first completed task
+        done, tasks = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+
+        if stop_event.is_set():
+            for task in tasks:
+                task.cancel()
+            return
+
+        # Update subscription of all finished tasks
+        for result in done:
+            domain, key = pending_lookup[result]
+            event_msg = {
+                "type": "update",
+                "domain": domain,
+                "key": key,
+            }
+
+            if result.exception():
+                exception = result.exception()
+                _LOGGER.error(
+                    "Error fetching system info for %s - %s",
+                    domain,
+                    key,
+                    exc_info=(type(exception), exception, exception.__traceback__),
+                )
+                event_msg["success"] = False
+                event_msg["error"] = {"type": "failed", "error": "unknown"}
+            else:
+                event_msg["success"] = True
+                event_msg["data"] = _format_value(result.result())
+
+            connection.send_message(
+                websocket_api.messages.event_message(msg["id"], event_msg)
+            )
+
+    connection.send_message(
+        websocket_api.messages.event_message(msg["id"], {"type": "finish"})
+    )
 
 
-@dataclasses.dataclass(frozen=True)
-class RegisterSystemHealth:
-    """Helper class to allow platforms to register."""
+@dataclasses.dataclass()
+class SystemHealthRegistration:
+    """Helper class to track platform registration."""
 
     hass: HomeAssistant
     domain: str
+    info_callback: Optional[Callable[[HomeAssistant], Awaitable[Dict]]] = None
+    manage_url: Optional[str] = None
 
     @callback
     def async_register_info(
         self,
-        info_callback: Callable[[HomeAssistant], Dict],
+        info_callback: Callable[[HomeAssistant], Awaitable[Dict]],
+        manage_url: Optional[str] = None,
     ):
         """Register an info callback."""
-        self.hass.data[DOMAIN]["info"][self.domain] = info_callback
+        self.info_callback = info_callback
+        self.manage_url = manage_url
+        self.hass.data[DOMAIN][self.domain] = self
+
+
+async def async_check_can_reach_url(
+    hass: HomeAssistant, url: str, more_info: Optional[str] = None
+) -> str:
+    """Test if the url can be reached."""
+    session = aiohttp_client.async_get_clientsession(hass)
+
+    try:
+        await session.get(url, timeout=5)
+        return "ok"
+    except aiohttp.ClientError:
+        data = {"type": "failed", "error": "unreachable"}
+        if more_info is not None:
+            data["more_info"] = more_info
+        return data

tests/common.py

@@ -965,7 +965,7 @@ async def flush_store(store):
 
 async def get_system_health_info(hass, domain):
     """Get system health info."""
-    return await hass.data["system_health"]["info"][domain](hass)
+    return await hass.data["system_health"][domain].info_callback(hass)
 
 
 def mock_integration(hass, module):

tests/components/cloud/test_system_health.py

@@ -0,0 +1,60 @@
+"""Test cloud system health."""
+import asyncio
+
+from aiohttp import ClientError
+
+from homeassistant.setup import async_setup_component
+from homeassistant.util.dt import utcnow
+
+from tests.async_mock import Mock
+from tests.common import get_system_health_info
+
+
+async def test_cloud_system_health(hass, aioclient_mock):
+    """Test cloud system health."""
+    aioclient_mock.get("https://cloud.bla.com/status", text="")
+    aioclient_mock.get("https://cert-server", text="")
+    aioclient_mock.get(
+        "https://cognito-idp.us-east-1.amazonaws.com/AAAA/.well-known/jwks.json",
+        exc=ClientError,
+    )
+    hass.config.components.add("cloud")
+    assert await async_setup_component(hass, "system_health", {})
+    now = utcnow()
+
+    hass.data["cloud"] = Mock(
+        region="us-east-1",
+        user_pool_id="AAAA",
+        relayer="wss://cloud.bla.com/websocket_api",
+        acme_directory_server="https://cert-server",
+        is_logged_in=True,
+        remote=Mock(is_connected=False),
+        expiration_date=now,
+        is_connected=True,
+        client=Mock(
+            prefs=Mock(
+                remote_enabled=True,
+                alexa_enabled=True,
+                google_enabled=False,
+            )
+        ),
+    )
+
+    info = await get_system_health_info(hass, "cloud")
+
+    for key, val in info.items():
+        if asyncio.iscoroutine(val):
+            info[key] = await val
+
+    assert info == {
+        "logged_in": True,
+        "subscription_expiration": now,
+        "relayer_connected": True,
+        "remote_enabled": True,
+        "remote_connected": False,
+        "alexa_enabled": True,
+        "google_enabled": False,
+        "can_reach_cert_server": "ok",
+        "can_reach_cloud_auth": {"type": "failed", "error": "unreachable"},
+        "can_reach_cloud": "ok",
+    }