Add subscription info endpoint (#16727)

* Add subscription info endpoint

* Lint

* Lint

* Make decorator

* Lint

homeassistant/components/cloud/__init__.py

@@ -23,7 +23,7 @@ from homeassistant.components.alexa import smart_home as alexa_sh
 from homeassistant.components.google_assistant import helpers as ga_h
 from homeassistant.components.google_assistant import const as ga_c
 
-from . import http_api, iot
+from . import http_api, iot, auth_api
 from .const import CONFIG_DIR, DOMAIN, SERVERS
 
 REQUIREMENTS = ['warrant==0.6.1']
@@ -39,6 +39,7 @@ CONF_GOOGLE_ACTIONS = 'google_actions'
 CONF_RELAYER = 'relayer'
 CONF_USER_POOL_ID = 'user_pool_id'
 CONF_GOOGLE_ACTIONS_SYNC_URL = 'google_actions_sync_url'
+CONF_SUBSCRIPTION_INFO_URL = 'subscription_info_url'
 
 DEFAULT_MODE = 'production'
 DEPENDENCIES = ['http']
@@ -79,6 +80,7 @@ CONFIG_SCHEMA = vol.Schema({
         vol.Optional(CONF_REGION): str,
         vol.Optional(CONF_RELAYER): str,
         vol.Optional(CONF_GOOGLE_ACTIONS_SYNC_URL): str,
+        vol.Optional(CONF_SUBSCRIPTION_INFO_URL): str,
         vol.Optional(CONF_ALEXA): ALEXA_SCHEMA,
         vol.Optional(CONF_GOOGLE_ACTIONS): GACTIONS_SCHEMA,
     }),
@@ -114,7 +116,8 @@ class Cloud:
 
     def __init__(self, hass, mode, alexa, google_actions,
                  cognito_client_id=None, user_pool_id=None, region=None,
-                 relayer=None, google_actions_sync_url=None):
+                 relayer=None, google_actions_sync_url=None,
+                 subscription_info_url=None):
         """Create an instance of Cloud."""
         self.hass = hass
         self.mode = mode
@@ -133,6 +136,7 @@ class Cloud:
             self.region = region
             self.relayer = relayer
             self.google_actions_sync_url = google_actions_sync_url
+            self.subscription_info_url = subscription_info_url
 
         else:
             info = SERVERS[mode]
@@ -142,6 +146,7 @@ class Cloud:
             self.region = info['region']
             self.relayer = info['relayer']
             self.google_actions_sync_url = info['google_actions_sync_url']
+            self.subscription_info_url = info['subscription_info_url']
 
     @property
     def is_logged_in(self):
@@ -195,6 +200,15 @@ class Cloud:
         """
         return self.hass.config.path(CONFIG_DIR, *parts)
 
+    async def fetch_subscription_info(self):
+        """Fetch subscription info."""
+        await self.hass.async_add_executor_job(auth_api.check_token, self)
+        websession = self.hass.helpers.aiohttp_client.async_get_clientsession()
+        return await websession.get(
+            self.subscription_info_url, headers={
+                'authorization': self.id_token
+            })
+
     @asyncio.coroutine
     def logout(self):
         """Close connection and remove all credentials."""

homeassistant/components/cloud/const.py

@@ -11,6 +11,8 @@ SERVERS = {
         'relayer': 'wss://cloud.hass.io:8000/websocket',
         'google_actions_sync_url': ('https://24ab3v80xd.execute-api.us-east-1.'
                                     'amazonaws.com/prod/smart_home_sync'),
+        'subscription_info_url': ('https://stripe-api.nabucasa.com/payments/'
+                                  'subscription_info')
     }
 }
 

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.
+        return self.json({'success': True})
-        await asyncio.sleep(0, loop=hass.loop)
-        return self.json(_account_data(cloud))
 
 
 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,
     }

homeassistant/components/websocket_api.py

@@ -480,6 +480,26 @@ class ActiveConnection:
         return wsock
 
 
+def async_response(func):
+    """Decorate an async function to handle WebSocket API messages."""
+    async def handle_msg_response(hass, connection, msg):
+        """Create a response and handle exception."""
+        try:
+            await func(hass, connection, msg)
+        except Exception:  # pylint: disable=broad-except
+            _LOGGER.exception("Unexpected exception")
+            connection.send_message_outside(error_message(
+                msg['id'], 'unknown', 'Unexpected error occurred'))
+
+    @callback
+    @wraps(func)
+    def schedule_handler(hass, connection, msg):
+        """Schedule the handler."""
+        hass.async_create_task(handle_msg_response(hass, connection, msg))
+
+    return schedule_handler
+
+
 @callback
 def handle_subscribe_events(hass, connection, msg):
     """Handle subscribe events command.
@@ -515,14 +535,12 @@ def handle_unsubscribe_events(hass, connection, msg):
             msg['id'], ERR_NOT_FOUND, 'Subscription not found.'))
 
 
-@callback
+@async_response
-def handle_call_service(hass, connection, msg):
+async def handle_call_service(hass, connection, msg):
     """Handle call service command.
 
     Async friendly.
     """
-    async def call_service_helper(msg):
-        """Call a service and fire complete message."""
     blocking = True
     if (msg['domain'] == 'homeassistant' and
             msg['service'] in ['restart', 'stop']):
@@ -532,8 +550,6 @@ def handle_call_service(hass, connection, msg):
         connection.context(msg))
     connection.send_message_outside(result_message(msg['id']))
 
-    hass.async_add_job(call_service_helper(msg))
-
 
 @callback
 def handle_get_states(hass, connection, msg):
@@ -545,20 +561,16 @@ def handle_get_states(hass, connection, msg):
         msg['id'], hass.states.async_all()))
 
 
-@callback
+@async_response
-def handle_get_services(hass, connection, msg):
+async def handle_get_services(hass, connection, msg):
     """Handle get services command.
 
     Async friendly.
     """
-    async def get_services_helper(msg):
-        """Get available services and fire complete message."""
     descriptions = await async_get_all_descriptions(hass)
     connection.send_message_outside(
         result_message(msg['id'], descriptions))
 
-    hass.async_add_job(get_services_helper(msg))
-
 
 @callback
 def handle_get_config(hass, connection, msg):

tests/components/cloud/test_http_api.py

@@ -12,14 +12,23 @@ from tests.common import mock_coro
 
 
 GOOGLE_ACTIONS_SYNC_URL = 'https://api-test.hass.io/google_actions_sync'
+SUBSCRIPTION_INFO_URL = 'https://api-test.hass.io/subscription_info'
 
 
-@pytest.fixture
+@pytest.fixture()
-def cloud_client(hass, aiohttp_client):
+def mock_auth():
-    """Fixture that can fetch from the cloud client."""
+    """Mock check token."""
+    with patch('homeassistant.components.cloud.auth_api.check_token'):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def setup_api(hass):
+    """Initialize HTTP API."""
     with patch('homeassistant.components.cloud.Cloud.async_start',
                return_value=mock_coro()):
-        hass.loop.run_until_complete(async_setup_component(hass, 'cloud', {
+        assert hass.loop.run_until_complete(async_setup_component(
+            hass, 'cloud', {
                 'cloud': {
                     'mode': 'development',
                     'cognito_client_id': 'cognito_client_id',
@@ -27,10 +36,16 @@ def cloud_client(hass, aiohttp_client):
                     'region': 'region',
                     'relayer': 'relayer',
                     'google_actions_sync_url': GOOGLE_ACTIONS_SYNC_URL,
+                    'subscription_info_url': SUBSCRIPTION_INFO_URL,
                 }
             }))
     hass.data['cloud']._decode_claims = \
         lambda token: jwt.get_unverified_claims(token)
+
+
+@pytest.fixture
+def cloud_client(hass, aiohttp_client):
+    """Fixture that can fetch from the cloud client."""
     with patch('homeassistant.components.cloud.Cloud.write_user_info'):
         yield hass.loop.run_until_complete(aiohttp_client(hass.http.app))
 
@@ -57,31 +72,6 @@ async def test_google_actions_sync_fails(mock_cognito, cloud_client,
     assert req.status == 403
 
 
-@asyncio.coroutine
-def test_account_view_no_account(cloud_client):
-    """Test fetching account if no account available."""
-    req = yield from cloud_client.get('/api/cloud/account')
-    assert req.status == 400
-
-
-@asyncio.coroutine
-def test_account_view(hass, cloud_client):
-    """Test fetching account if no account available."""
-    hass.data[DOMAIN].id_token = jwt.encode({
-        'email': 'hello@home-assistant.io',
-        'custom:sub-exp': '2018-01-03'
-    }, 'test')
-    hass.data[DOMAIN].iot.state = iot.STATE_CONNECTED
-    req = yield from cloud_client.get('/api/cloud/account')
-    assert req.status == 200
-    result = yield from req.json()
-    assert result == {
-        'email': 'hello@home-assistant.io',
-        'sub_exp': '2018-01-03',
-        'cloud': iot.STATE_CONNECTED,
-    }
-
-
 @asyncio.coroutine
 def test_login_view(hass, cloud_client, mock_cognito):
     """Test logging in."""
@@ -103,8 +93,7 @@ def test_login_view(hass, cloud_client, mock_cognito):
 
     assert req.status == 200
     result = yield from req.json()
-    assert result['email'] == 'hello@home-assistant.io'
+    assert result == {'success': True}
-    assert result['sub_exp'] == '2018-01-03'
 
     assert len(mock_connect.mock_calls) == 1
 
@@ -330,3 +319,91 @@ def test_resend_confirm_view_unknown_error(mock_cognito, cloud_client):
         'email': 'hello@bla.com',
     })
     assert req.status == 502
+
+
+async def test_websocket_status(hass, hass_ws_client):
+    """Test querying the status."""
+    hass.data[DOMAIN].id_token = jwt.encode({
+        'email': 'hello@home-assistant.io',
+        'custom:sub-exp': '2018-01-03'
+    }, 'test')
+    hass.data[DOMAIN].iot.state = iot.STATE_CONNECTED
+    client = await hass_ws_client(hass)
+    await client.send_json({
+        'id': 5,
+        'type': 'cloud/status'
+    })
+    response = await client.receive_json()
+    assert response['result'] == {
+        'logged_in': True,
+        'email': 'hello@home-assistant.io',
+        'cloud': 'connected',
+    }
+
+
+async def test_websocket_status_not_logged_in(hass, hass_ws_client):
+    """Test querying the status."""
+    client = await hass_ws_client(hass)
+    await client.send_json({
+        'id': 5,
+        'type': 'cloud/status'
+    })
+    response = await client.receive_json()
+    assert response['result'] == {
+        'logged_in': False,
+        'cloud': 'disconnected'
+    }
+
+
+async def test_websocket_subscription(hass, hass_ws_client, aioclient_mock,
+                                      mock_auth):
+    """Test querying the status."""
+    aioclient_mock.get(SUBSCRIPTION_INFO_URL, json={'return': 'value'})
+    hass.data[DOMAIN].id_token = jwt.encode({
+        'email': 'hello@home-assistant.io',
+        'custom:sub-exp': '2018-01-03'
+    }, 'test')
+    client = await hass_ws_client(hass)
+    await client.send_json({
+        'id': 5,
+        'type': 'cloud/subscription'
+    })
+    response = await client.receive_json()
+
+    assert response['result'] == {
+        'return': 'value'
+    }
+
+
+async def test_websocket_subscription_fail(hass, hass_ws_client,
+                                           aioclient_mock, mock_auth):
+    """Test querying the status."""
+    aioclient_mock.get(SUBSCRIPTION_INFO_URL, status=500)
+    hass.data[DOMAIN].id_token = jwt.encode({
+        'email': 'hello@home-assistant.io',
+        'custom:sub-exp': '2018-01-03'
+    }, 'test')
+    client = await hass_ws_client(hass)
+    await client.send_json({
+        'id': 5,
+        'type': 'cloud/subscription'
+    })
+    response = await client.receive_json()
+
+    assert not response['success']
+    assert response['error']['code'] == 'request_failed'
+
+
+async def test_websocket_subscription_not_logged_in(hass, hass_ws_client):
+    """Test querying the status."""
+    client = await hass_ws_client(hass)
+    with patch('homeassistant.components.cloud.Cloud.fetch_subscription_info',
+               return_value=mock_coro({'return': 'value'})):
+        await client.send_json({
+            'id': 5,
+            'type': 'cloud/subscription'
+        })
+        response = await client.receive_json()
+
+    assert not response['success']
+    assert response['error']['code'] == 'not_logged_in'

tests/components/cloud/test_init.py

@@ -30,6 +30,7 @@ def test_constructor_loads_info_from_constant():
             'region': 'test-region',
             'relayer': 'test-relayer',
             'google_actions_sync_url': 'test-google_actions_sync_url',
+            'subscription_info_url': 'test-subscription-info-url'
         }
     }), patch('homeassistant.components.cloud.Cloud._fetch_jwt_keyset',
               return_value=mock_coro(True)):
@@ -45,6 +46,7 @@ def test_constructor_loads_info_from_constant():
     assert cl.region == 'test-region'
     assert cl.relayer == 'test-relayer'
     assert cl.google_actions_sync_url == 'test-google_actions_sync_url'
+    assert cl.subscription_info_url == 'test-subscription-info-url'
 
 
 @asyncio.coroutine