Add subscription info endpoint (#16727)

* Add subscription info endpoint

* Lint

* Lint

* Make decorator

* Lint

homeassistant/components/cloud/http_api.py

@@ -6,22 +6,44 @@ import logging
 import async_timeout
 import voluptuous as vol
 
+from homeassistant.core import callback
 from homeassistant.components.http import HomeAssistantView
 from homeassistant.components.http.data_validator import (
     RequestDataValidator)
+from homeassistant.components import websocket_api
 
 from . import auth_api
 from .const import DOMAIN, REQUEST_TIMEOUT
+from .iot import STATE_DISCONNECTED
 
 _LOGGER = logging.getLogger(__name__)
 
 
+WS_TYPE_STATUS = 'cloud/status'
+SCHEMA_WS_STATUS = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_STATUS,
+})
+
+
+WS_TYPE_SUBSCRIPTION = 'cloud/subscription'
+SCHEMA_WS_SUBSCRIPTION = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
+    vol.Required('type'): WS_TYPE_SUBSCRIPTION,
+})
+
+
 async def async_setup(hass):
     """Initialize the HTTP API."""
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_STATUS, websocket_cloud_status,
+        SCHEMA_WS_STATUS
+    )
+    hass.components.websocket_api.async_register_command(
+        WS_TYPE_SUBSCRIPTION, websocket_subscription,
+        SCHEMA_WS_SUBSCRIPTION
+    )
     hass.http.register_view(GoogleActionsSyncView)
     hass.http.register_view(CloudLoginView)
     hass.http.register_view(CloudLogoutView)
-    hass.http.register_view(CloudAccountView)
     hass.http.register_view(CloudRegisterView)
     hass.http.register_view(CloudResendConfirmView)
     hass.http.register_view(CloudForgotPasswordView)
@@ -102,9 +124,7 @@ class CloudLoginView(HomeAssistantView):
                                      data['password'])
 
         hass.async_add_job(cloud.iot.connect)
-        # Allow cloud to start connecting.
-        await asyncio.sleep(0, loop=hass.loop)
-        return self.json(_account_data(cloud))
+        return self.json({'success': True})
 
 
 class CloudLogoutView(HomeAssistantView):
@@ -125,23 +145,6 @@ class CloudLogoutView(HomeAssistantView):
         return self.json_message('ok')
 
 
-class CloudAccountView(HomeAssistantView):
-    """View to retrieve account info."""
-
-    url = '/api/cloud/account'
-    name = 'api:cloud:account'
-
-    async def get(self, request):
-        """Get account info."""
-        hass = request.app['hass']
-        cloud = hass.data[DOMAIN]
-
-        if not cloud.is_logged_in:
-            return self.json_message('Not logged in', 400)
-
-        return self.json(_account_data(cloud))
-
-
 class CloudRegisterView(HomeAssistantView):
     """Register on the Home Assistant cloud."""
 
@@ -209,12 +212,51 @@ class CloudForgotPasswordView(HomeAssistantView):
         return self.json_message('ok')
 
 
+@callback
+def websocket_cloud_status(hass, connection, msg):
+    """Handle request for account info.
+
+    Async friendly.
+    """
+    cloud = hass.data[DOMAIN]
+    connection.to_write.put_nowait(
+        websocket_api.result_message(msg['id'], _account_data(cloud)))
+
+
+@websocket_api.async_response
+async def websocket_subscription(hass, connection, msg):
+    """Handle request for account info."""
+    cloud = hass.data[DOMAIN]
+
+    if not cloud.is_logged_in:
+        connection.to_write.put_nowait(websocket_api.error_message(
+            msg['id'], 'not_logged_in',
+            'You need to be logged in to the cloud.'))
+        return
+
+    with async_timeout.timeout(REQUEST_TIMEOUT, loop=hass.loop):
+        response = await cloud.fetch_subscription_info()
+
+    if response.status == 200:
+        connection.send_message_outside(websocket_api.result_message(
+            msg['id'], await response.json()))
+    else:
+        connection.send_message_outside(websocket_api.error_message(
+            msg['id'], 'request_failed', 'Failed to request subscription'))
+
+
 def _account_data(cloud):
     """Generate the auth data JSON response."""
+    if not cloud.is_logged_in:
+        return {
+            'logged_in': False,
+            'cloud': STATE_DISCONNECTED,
+        }
+
     claims = cloud.claims
 
     return {
+        'logged_in': True,
         'email': claims['email'],
-        'sub_exp': claims['custom:sub-exp'],
         'cloud': cloud.iot.state,
     }