Cloud connection via aiohttp (#9860)

* Cloud: connect to cloud

* Fix tests in py34

* Update warrant to 0.5.0

* Differentiate errors between unknown handler vs exception

* Lint

* Respond to cloud message to logout

* Refresh token exception handling

* Swap out bare exception for RuntimeError

* Add more tests

* Fix tests py34

homeassistant/components/cloud/__init__.py

@@ -1,47 +1,147 @@
 """Component to integrate the Home Assistant cloud."""
 import asyncio
+import json
 import logging
+import os
 
 import voluptuous as vol
 
-from . import http_api, auth_api
+from homeassistant.const import EVENT_HOMEASSISTANT_START
-from .const import DOMAIN
+
+from . import http_api, iot
+from .const import CONFIG_DIR, DOMAIN, SERVERS
 
 
-REQUIREMENTS = ['warrant==0.2.0']
+REQUIREMENTS = ['warrant==0.5.0']
 DEPENDENCIES = ['http']
 CONF_MODE = 'mode'
+CONF_COGNITO_CLIENT_ID = 'cognito_client_id'
+CONF_USER_POOL_ID = 'user_pool_id'
+CONF_REGION = 'region'
+CONF_RELAYER = 'relayer'
 MODE_DEV = 'development'
-MODE_STAGING = 'staging'
-MODE_PRODUCTION = 'production'
 DEFAULT_MODE = MODE_DEV
+_LOGGER = logging.getLogger(__name__)
 
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
         vol.Optional(CONF_MODE, default=DEFAULT_MODE):
-        vol.In([MODE_DEV, MODE_STAGING, MODE_PRODUCTION]),
+            vol.In([MODE_DEV] + list(SERVERS)),
+        # Change to optional when we include real servers
+        vol.Required(CONF_COGNITO_CLIENT_ID): str,
+        vol.Required(CONF_USER_POOL_ID): str,
+        vol.Required(CONF_REGION): str,
+        vol.Required(CONF_RELAYER): str,
     }),
 }, extra=vol.ALLOW_EXTRA)
-_LOGGER = logging.getLogger(__name__)
 
 
 @asyncio.coroutine
 def async_setup(hass, config):
     """Initialize the Home Assistant cloud."""
-    mode = MODE_PRODUCTION
-
     if DOMAIN in config:
-        mode = config[DOMAIN].get(CONF_MODE)
+        kwargs = config[DOMAIN]
+    else:
+        kwargs = {CONF_MODE: DEFAULT_MODE}
 
-    if mode != 'development':
+    cloud = hass.data[DOMAIN] = Cloud(hass, **kwargs)
-        _LOGGER.error('Only development mode is currently allowed.')
-        return False
 
-    data = hass.data[DOMAIN] = {
+    @asyncio.coroutine
-        'mode': mode
+    def init_cloud(event):
-    }
+        """Initialize connection."""
+        yield from cloud.initialize()
 
-    data['auth'] = yield from hass.async_add_job(auth_api.load_auth, hass)
+    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START, init_cloud)
 
     yield from http_api.async_setup(hass)
     return True
+
+
+class Cloud:
+    """Store the configuration of the cloud connection."""
+
+    def __init__(self, hass, mode, cognito_client_id=None, user_pool_id=None,
+                 region=None, relayer=None):
+        """Create an instance of Cloud."""
+        self.hass = hass
+        self.mode = mode
+        self.email = None
+        self.id_token = None
+        self.access_token = None
+        self.refresh_token = None
+        self.iot = iot.CloudIoT(self)
+
+        if mode == MODE_DEV:
+            self.cognito_client_id = cognito_client_id
+            self.user_pool_id = user_pool_id
+            self.region = region
+            self.relayer = relayer
+
+        else:
+            info = SERVERS[mode]
+
+            self.cognito_client_id = info['cognito_client_id']
+            self.user_pool_id = info['user_pool_id']
+            self.region = info['region']
+            self.relayer = info['relayer']
+
+    @property
+    def is_logged_in(self):
+        """Get if cloud is logged in."""
+        return self.email is not None
+
+    @property
+    def user_info_path(self):
+        """Get path to the stored auth."""
+        return self.path('{}_auth.json'.format(self.mode))
+
+    @asyncio.coroutine
+    def initialize(self):
+        """Initialize and load cloud info."""
+        def load_config():
+            """Load the configuration."""
+            # Ensure config dir exists
+            path = self.hass.config.path(CONFIG_DIR)
+            if not os.path.isdir(path):
+                os.mkdir(path)
+
+            user_info = self.user_info_path
+            if os.path.isfile(user_info):
+                with open(user_info, 'rt') as file:
+                    info = json.loads(file.read())
+                self.email = info['email']
+                self.id_token = info['id_token']
+                self.access_token = info['access_token']
+                self.refresh_token = info['refresh_token']
+
+        yield from self.hass.async_add_job(load_config)
+
+        if self.email is not None:
+            yield from self.iot.connect()
+
+    def path(self, *parts):
+        """Get config path inside cloud dir."""
+        return self.hass.config.path(CONFIG_DIR, *parts)
+
+    @asyncio.coroutine
+    def logout(self):
+        """Close connection and remove all credentials."""
+        yield from self.iot.disconnect()
+
+        self.email = None
+        self.id_token = None
+        self.access_token = None
+        self.refresh_token = None
+
+        yield from self.hass.async_add_job(
+            lambda: os.remove(self.user_info_path))
+
+    def write_user_info(self):
+        """Write user info to a file."""
+        with open(self.user_info_path, 'wt') as file:
+            file.write(json.dumps({
+                'email': self.email,
+                'id_token': self.id_token,
+                'access_token': self.access_token,
+                'refresh_token': self.refresh_token,
+            }, indent=4))

homeassistant/components/cloud/auth_api.py

@@ -1,10 +1,7 @@
-"""Package to offer tools to authenticate with the cloud."""
+"""Package to communicate with the authentication API."""
-import json
+import hashlib
 import logging
-import os
 
-from .const import AUTH_FILE, SERVERS
-from .util import get_mode
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -61,116 +58,95 @@ def _map_aws_exception(err):
     return ex(err.response['Error']['Message'])
 
 
-def load_auth(hass):
+def _generate_username(email):
-    """Load authentication from disk and verify it."""
+    """Generate a username from an email address."""
-    info = _read_info(hass)
+    return hashlib.sha512(email.encode('utf-8')).hexdigest()
-
-    if info is None:
-        return Auth(hass)
-
-    auth = Auth(hass, _cognito(
-        hass,
-        id_token=info['id_token'],
-        access_token=info['access_token'],
-        refresh_token=info['refresh_token'],
-    ))
-
-    if auth.validate_auth():
-        return auth
-
-    return Auth(hass)
 
 
-def register(hass, email, password):
+def register(cloud, email, password):
     """Register a new account."""
     from botocore.exceptions import ClientError
 
-    cognito = _cognito(hass, username=email)
+    cognito = _cognito(cloud)
     try:
-        cognito.register(email, password)
+        cognito.register(_generate_username(email), password, email=email)
     except ClientError as err:
         raise _map_aws_exception(err)
 
 
-def confirm_register(hass, confirmation_code, email):
+def confirm_register(cloud, confirmation_code, email):
     """Confirm confirmation code after registration."""
     from botocore.exceptions import ClientError
 
-    cognito = _cognito(hass, username=email)
+    cognito = _cognito(cloud)
     try:
-        cognito.confirm_sign_up(confirmation_code, email)
+        cognito.confirm_sign_up(confirmation_code, _generate_username(email))
     except ClientError as err:
         raise _map_aws_exception(err)
 
 
-def forgot_password(hass, email):
+def forgot_password(cloud, email):
     """Initiate forgotten password flow."""
     from botocore.exceptions import ClientError
 
-    cognito = _cognito(hass, username=email)
+    cognito = _cognito(cloud, username=_generate_username(email))
     try:
         cognito.initiate_forgot_password()
     except ClientError as err:
         raise _map_aws_exception(err)
 
 
-def confirm_forgot_password(hass, confirmation_code, email, new_password):
+def confirm_forgot_password(cloud, confirmation_code, email, new_password):
     """Confirm forgotten password code and change password."""
     from botocore.exceptions import ClientError
 
-    cognito = _cognito(hass, username=email)
+    cognito = _cognito(cloud, username=_generate_username(email))
     try:
         cognito.confirm_forgot_password(confirmation_code, new_password)
     except ClientError as err:
         raise _map_aws_exception(err)
 
 
-class Auth(object):
+def login(cloud, email, password):
-    """Class that holds Cloud authentication."""
+    """Log user in and fetch certificate."""
+    cognito = _authenticate(cloud, email, password)
+    cloud.id_token = cognito.id_token
+    cloud.access_token = cognito.access_token
+    cloud.refresh_token = cognito.refresh_token
+    cloud.email = email
+    cloud.write_user_info()
 
-    def __init__(self, hass, cognito=None):
-        """Initialize Hass cloud info object."""
-        self.hass = hass
-        self.cognito = cognito
-        self.account = None
 
-    @property
+def check_token(cloud):
-    def is_logged_in(self):
+    """Check that the token is valid and verify if needed."""
-        """Return if user is logged in."""
-        return self.account is not None
-
-    def validate_auth(self):
-        """Validate that the contained auth is valid."""
     from botocore.exceptions import ClientError
 
+    cognito = _cognito(
+        cloud,
+        access_token=cloud.access_token,
+        refresh_token=cloud.refresh_token)
+
     try:
-            self._refresh_account_info()
+        if cognito.check_token():
+            cloud.id_token = cognito.id_token
+            cloud.access_token = cognito.access_token
+            cloud.write_user_info()
     except ClientError as err:
-            if err.response['Error']['Code'] != 'NotAuthorizedException':
+        raise _map_aws_exception(err)
-                _LOGGER.error('Unexpected error verifying auth: %s', err)
-                return False
 
-            try:
-                self.renew_access_token()
-                self._refresh_account_info()
-            except ClientError:
-                _LOGGER.error('Unable to refresh auth token: %s', err)
-                return False
 
-        return True
+def _authenticate(cloud, email, password):
-
+    """Log in and return an authenticated Cognito instance."""
-    def login(self, username, password):
-        """Login using a username and password."""
     from botocore.exceptions import ClientError
     from warrant.exceptions import ForceChangePasswordException
 
-        cognito = _cognito(self.hass, username=username)
+    assert not cloud.is_logged_in, 'Cannot login if already logged in.'
+
+    cognito = _cognito(cloud, username=email)
 
     try:
         cognito.authenticate(password=password)
-            self.cognito = cognito
+        return cognito
-            self._refresh_account_info()
-            _write_info(self.hass, self)
 
     except ForceChangePasswordException as err:
         raise PasswordChangeRequired
@@ -178,93 +154,24 @@ class Auth(object):
     except ClientError as err:
         raise _map_aws_exception(err)
 
-    def _refresh_account_info(self):
-        """Refresh the account info.
 
-        Raises boto3 exceptions.
+def _cognito(cloud, **kwargs):
-        """
-        self.account = self.cognito.get_user()
-
-    def renew_access_token(self):
-        """Refresh token."""
-        from botocore.exceptions import ClientError
-
-        try:
-            self.cognito.renew_access_token()
-            _write_info(self.hass, self)
-            return True
-        except ClientError as err:
-            _LOGGER.error('Error refreshing token: %s', err)
-            return False
-
-    def logout(self):
-        """Invalidate token."""
-        from botocore.exceptions import ClientError
-
-        try:
-            self.cognito.logout()
-            self.account = None
-            _write_info(self.hass, self)
-        except ClientError as err:
-            raise _map_aws_exception(err)
-
-
-def _read_info(hass):
-    """Read auth file."""
-    path = hass.config.path(AUTH_FILE)
-
-    if not os.path.isfile(path):
-        return None
-
-    with open(path) as file:
-        return json.load(file).get(get_mode(hass))
-
-
-def _write_info(hass, auth):
-    """Write auth info for specified mode.
-
-    Pass in None for data to remove authentication for that mode.
-    """
-    path = hass.config.path(AUTH_FILE)
-    mode = get_mode(hass)
-
-    if os.path.isfile(path):
-        with open(path) as file:
-            content = json.load(file)
-    else:
-        content = {}
-
-    if auth.is_logged_in:
-        content[mode] = {
-            'id_token': auth.cognito.id_token,
-            'access_token': auth.cognito.access_token,
-            'refresh_token': auth.cognito.refresh_token,
-        }
-    else:
-        content.pop(mode, None)
-
-    with open(path, 'wt') as file:
-        file.write(json.dumps(content, indent=4, sort_keys=True))
-
-
-def _cognito(hass, **kwargs):
     """Get the client credentials."""
+    import botocore
+    import boto3
     from warrant import Cognito
 
-    mode = get_mode(hass)
-
-    info = SERVERS.get(mode)
-
-    if info is None:
-        raise ValueError('Mode {} is not supported.'.format(mode))
-
     cognito = Cognito(
-        user_pool_id=info['identity_pool_id'],
+        user_pool_id=cloud.user_pool_id,
-        client_id=info['client_id'],
+        client_id=cloud.cognito_client_id,
-        user_pool_region=info['region'],
+        user_pool_region=cloud.region,
-        access_key=info['access_key_id'],
-        secret_key=info['secret_access_key'],
         **kwargs
     )
-
+    cognito.client = boto3.client(
+        'cognito-idp',
+        region_name=cloud.region,
+        config=botocore.config.Config(
+            signature_version=botocore.UNSIGNED
+        )
+    )
     return cognito

homeassistant/components/cloud/const.py

@@ -1,14 +1,14 @@
 """Constants for the cloud component."""
 DOMAIN = 'cloud'
+CONFIG_DIR = '.cloud'
 REQUEST_TIMEOUT = 10
-AUTH_FILE = '.cloud'
 
 SERVERS = {
-    'development': {
+    # Example entry:
-        'client_id': '3k755iqfcgv8t12o4pl662mnos',
+    # 'production': {
-        'identity_pool_id': 'us-west-2_vDOfweDJo',
+    #     'cognito_client_id': '',
-        'region': 'us-west-2',
+    #     'user_pool_id': '',
-        'access_key_id': 'AKIAJGRK7MILPRJTT2ZQ',
+    #     'region': '',
-        'secret_access_key': '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'
+    #     'relayer': ''
-    }
+    # }
 }

homeassistant/components/cloud/http_api.py

@@ -10,7 +10,7 @@ from homeassistant.components.http import (
     HomeAssistantView, RequestDataValidator)
 
 from . import auth_api
-from .const import REQUEST_TIMEOUT
+from .const import DOMAIN, REQUEST_TIMEOUT
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -74,13 +74,14 @@ class CloudLoginView(HomeAssistantView):
     def post(self, request, data):
         """Handle login request."""
         hass = request.app['hass']
-        auth = hass.data['cloud']['auth']
+        cloud = hass.data[DOMAIN]
 
         with async_timeout.timeout(REQUEST_TIMEOUT, loop=hass.loop):
-            yield from hass.async_add_job(auth.login, data['email'],
+            yield from hass.async_add_job(auth_api.login, cloud, data['email'],
                                           data['password'])
+            hass.async_add_job(cloud.iot.connect)
 
-        return self.json(_auth_data(auth))
+        return self.json(_account_data(cloud))
 
 
 class CloudLogoutView(HomeAssistantView):
@@ -94,10 +95,10 @@ class CloudLogoutView(HomeAssistantView):
     def post(self, request):
         """Handle logout request."""
         hass = request.app['hass']
-        auth = hass.data['cloud']['auth']
+        cloud = hass.data[DOMAIN]
 
         with async_timeout.timeout(REQUEST_TIMEOUT, loop=hass.loop):
-            yield from hass.async_add_job(auth.logout)
+            yield from cloud.logout()
 
         return self.json_message('ok')
 
@@ -112,12 +113,12 @@ class CloudAccountView(HomeAssistantView):
     def get(self, request):
         """Get account info."""
         hass = request.app['hass']
-        auth = hass.data['cloud']['auth']
+        cloud = hass.data[DOMAIN]
 
-        if not auth.is_logged_in:
+        if not cloud.is_logged_in:
             return self.json_message('Not logged in', 400)
 
-        return self.json(_auth_data(auth))
+        return self.json(_account_data(cloud))
 
 
 class CloudRegisterView(HomeAssistantView):
@@ -135,10 +136,11 @@ class CloudRegisterView(HomeAssistantView):
     def post(self, request, data):
         """Handle registration request."""
         hass = request.app['hass']
+        cloud = hass.data[DOMAIN]
 
         with async_timeout.timeout(REQUEST_TIMEOUT, loop=hass.loop):
             yield from hass.async_add_job(
-                auth_api.register, hass, data['email'], data['password'])
+                auth_api.register, cloud, data['email'], data['password'])
 
         return self.json_message('ok')
 
@@ -158,10 +160,11 @@ class CloudConfirmRegisterView(HomeAssistantView):
     def post(self, request, data):
         """Handle registration confirmation request."""
         hass = request.app['hass']
+        cloud = hass.data[DOMAIN]
 
         with async_timeout.timeout(REQUEST_TIMEOUT, loop=hass.loop):
             yield from hass.async_add_job(
-                auth_api.confirm_register, hass, data['confirmation_code'],
+                auth_api.confirm_register, cloud, data['confirmation_code'],
                 data['email'])
 
         return self.json_message('ok')
@@ -181,10 +184,11 @@ class CloudForgotPasswordView(HomeAssistantView):
     def post(self, request, data):
         """Handle forgot password request."""
         hass = request.app['hass']
+        cloud = hass.data[DOMAIN]
 
         with async_timeout.timeout(REQUEST_TIMEOUT, loop=hass.loop):
             yield from hass.async_add_job(
-                auth_api.forgot_password, hass, data['email'])
+                auth_api.forgot_password, cloud, data['email'])
 
         return self.json_message('ok')
 
@@ -205,18 +209,19 @@ class CloudConfirmForgotPasswordView(HomeAssistantView):
     def post(self, request, data):
         """Handle forgot password confirm request."""
         hass = request.app['hass']
+        cloud = hass.data[DOMAIN]
 
         with async_timeout.timeout(REQUEST_TIMEOUT, loop=hass.loop):
             yield from hass.async_add_job(
-                auth_api.confirm_forgot_password, hass,
+                auth_api.confirm_forgot_password, cloud,
                 data['confirmation_code'], data['email'],
                 data['new_password'])
 
         return self.json_message('ok')
 
 
-def _auth_data(auth):
+def _account_data(cloud):
     """Generate the auth data JSON response."""
     return {
-        'email': auth.account.email
+        'email': cloud.email
     }

homeassistant/components/cloud/iot.py

@@ -0,0 +1,195 @@
+"""Module to handle messages from Home Assistant cloud."""
+import asyncio
+import logging
+
+from aiohttp import hdrs, client_exceptions, WSMsgType
+
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.components.alexa import smart_home
+from homeassistant.util.decorator import Registry
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from . import auth_api
+
+
+HANDLERS = Registry()
+_LOGGER = logging.getLogger(__name__)
+
+
+class UnknownHandler(Exception):
+    """Exception raised when trying to handle unknown handler."""
+
+
+class CloudIoT:
+    """Class to manage the IoT connection."""
+
+    def __init__(self, cloud):
+        """Initialize the CloudIoT class."""
+        self.cloud = cloud
+        self.client = None
+        self.close_requested = False
+        self.tries = 0
+
+    @property
+    def is_connected(self):
+        """Return if connected to the cloud."""
+        return self.client is not None
+
+    @asyncio.coroutine
+    def connect(self):
+        """Connect to the IoT broker."""
+        if self.client is not None:
+            raise RuntimeError('Cannot connect while already connected')
+
+        self.close_requested = False
+
+        hass = self.cloud.hass
+        remove_hass_stop_listener = None
+
+        session = async_get_clientsession(self.cloud.hass)
+        headers = {
+            hdrs.AUTHORIZATION: 'Bearer {}'.format(self.cloud.access_token)
+        }
+
+        @asyncio.coroutine
+        def _handle_hass_stop(event):
+            """Handle Home Assistant shutting down."""
+            nonlocal remove_hass_stop_listener
+            remove_hass_stop_listener = None
+            yield from self.disconnect()
+
+        client = None
+        disconnect_warn = None
+        try:
+            yield from hass.async_add_job(auth_api.check_token, self.cloud)
+
+            self.client = client = yield from session.ws_connect(
+                'ws://{}/websocket'.format(self.cloud.relayer),
+                headers=headers)
+            self.tries = 0
+
+            remove_hass_stop_listener = hass.bus.async_listen_once(
+                EVENT_HOMEASSISTANT_STOP, _handle_hass_stop)
+
+            _LOGGER.info('Connected')
+
+            while not client.closed:
+                msg = yield from client.receive()
+
+                if msg.type in (WSMsgType.ERROR, WSMsgType.CLOSED,
+                                WSMsgType.CLOSING):
+                    disconnect_warn = 'Closed by server'
+                    break
+
+                elif msg.type != WSMsgType.TEXT:
+                    disconnect_warn = 'Received non-Text message: {}'.format(
+                        msg.type)
+                    break
+
+                try:
+                    msg = msg.json()
+                except ValueError:
+                    disconnect_warn = 'Received invalid JSON.'
+                    break
+
+                _LOGGER.debug('Received message: %s', msg)
+
+                response = {
+                    'msgid': msg['msgid'],
+                }
+                try:
+                    result = yield from async_handle_message(
+                        hass, self.cloud, msg['handler'], msg['payload'])
+
+                    # No response from handler
+                    if result is None:
+                        continue
+
+                    response['payload'] = result
+
+                except UnknownHandler:
+                    response['error'] = 'unknown-handler'
+
+                except Exception:  # pylint: disable=broad-except
+                    _LOGGER.exception('Error handling message')
+                    response['error'] = 'exception'
+
+                _LOGGER.debug('Publishing message: %s', response)
+                yield from client.send_json(response)
+
+        except auth_api.CloudError:
+            _LOGGER.warning('Unable to connect: Unable to refresh token.')
+
+        except client_exceptions.WSServerHandshakeError as err:
+            if err.code == 401:
+                disconnect_warn = 'Invalid auth.'
+                self.close_requested = True
+                # Should we notify user?
+            else:
+                _LOGGER.warning('Unable to connect: %s', err)
+
+        except client_exceptions.ClientError as err:
+            _LOGGER.warning('Unable to connect: %s', err)
+
+        except Exception:  # pylint: disable=broad-except
+            if not self.close_requested:
+                _LOGGER.exception('Unexpected error')
+
+        finally:
+            if disconnect_warn is not None:
+                _LOGGER.warning('Connection closed: %s', disconnect_warn)
+
+            if remove_hass_stop_listener is not None:
+                remove_hass_stop_listener()
+
+            if client is not None:
+                self.client = None
+                yield from client.close()
+
+            if not self.close_requested:
+                self.tries += 1
+
+                # Sleep 0, 5, 10, 15 … up to 30 seconds between retries
+                yield from asyncio.sleep(
+                    min(30, (self.tries - 1) * 5), loop=hass.loop)
+
+                hass.async_add_job(self.connect())
+
+    @asyncio.coroutine
+    def disconnect(self):
+        """Disconnect the client."""
+        self.close_requested = True
+        yield from self.client.close()
+
+
+@asyncio.coroutine
+def async_handle_message(hass, cloud, handler_name, payload):
+    """Handle incoming IoT message."""
+    handler = HANDLERS.get(handler_name)
+
+    if handler is None:
+        raise UnknownHandler()
+
+    return (yield from handler(hass, cloud, payload))
+
+
+@HANDLERS.register('alexa')
+@asyncio.coroutine
+def async_handle_alexa(hass, cloud, payload):
+    """Handle an incoming IoT message for Alexa."""
+    return (yield from smart_home.async_handle_message(hass, payload))
+
+
+@HANDLERS.register('cloud')
+@asyncio.coroutine
+def async_handle_cloud(hass, cloud, payload):
+    """Handle an incoming IoT message for cloud component."""
+    action = payload['action']
+
+    if action == 'logout':
+        yield from cloud.logout()
+        _LOGGER.error('You have been logged out from Home Assistant cloud: %s',
+                      payload['reason'])
+    else:
+        _LOGGER.warning('Received unknown cloud action: %s', action)
+
+    return None

homeassistant/components/cloud/util.py

@@ -1,10 +0,0 @@
-"""Utilities for the cloud integration."""
-from .const import DOMAIN
-
-
-def get_mode(hass):
-    """Return the current mode of the cloud component.
-
-    Async friendly.
-    """
-    return hass.data[DOMAIN]['mode']

requirements_all.txt

@@ -1043,7 +1043,7 @@ wakeonlan==0.2.2
 waqiasync==1.0.0
 
 # homeassistant.components.cloud
-warrant==0.2.0
+warrant==0.5.0
 
 # homeassistant.components.media_player.gpmdp
 websocket-client==0.37.0

requirements_test_all.txt

@@ -149,7 +149,7 @@ statsd==3.2.1
 uvcclient==0.10.1
 
 # homeassistant.components.cloud
-warrant==0.2.0
+warrant==0.5.0
 
 # homeassistant.components.sensor.yahoo_finance
 yahoo-finance==1.4.0

tests/components/cloud/test_auth_api.py

@@ -4,35 +4,7 @@ from unittest.mock import MagicMock, patch
 from botocore.exceptions import ClientError
 import pytest
 
-from homeassistant.components.cloud import DOMAIN, auth_api
+from homeassistant.components.cloud import auth_api
-
-
-MOCK_AUTH = {
-    "id_token": "fake_id_token",
-    "access_token": "fake_access_token",
-    "refresh_token": "fake_refresh_token",
-}
-
-
-@pytest.fixture
-def cloud_hass(hass):
-    """Fixture to return a hass instance with cloud mode set."""
-    hass.data[DOMAIN] = {'mode': 'development'}
-    return hass
-
-
-@pytest.fixture
-def mock_write():
-    """Mock reading authentication."""
-    with patch.object(auth_api, '_write_info') as mock:
-        yield mock
-
-
-@pytest.fixture
-def mock_read():
-    """Mock writing authentication."""
-    with patch.object(auth_api, '_read_info') as mock:
-        yield mock
 
 
 @pytest.fixture
@@ -42,13 +14,6 @@ def mock_cognito():
         yield mock_cog()
 
 
-@pytest.fixture
-def mock_auth():
-    """Mock warrant."""
-    with patch('homeassistant.components.cloud.auth_api.Auth') as mock_auth:
-        yield mock_auth()
-
-
 def aws_error(code, message='Unknown', operation_name='fake_operation_name'):
     """Generate AWS error response."""
     response = {
@@ -60,159 +25,64 @@ def aws_error(code, message='Unknown', operation_name='fake_operation_name'):
     return ClientError(response, operation_name)
 
 
-def test_load_auth_with_no_stored_auth(cloud_hass, mock_read):
+def test_login_invalid_auth(mock_cognito):
-    """Test loading authentication with no stored auth."""
-    mock_read.return_value = None
-    auth = auth_api.load_auth(cloud_hass)
-    assert auth.cognito is None
-
-
-def test_load_auth_with_invalid_auth(cloud_hass, mock_read, mock_cognito):
-    """Test calling load_auth when auth is no longer valid."""
-    mock_cognito.get_user.side_effect = aws_error('SomeError')
-    auth = auth_api.load_auth(cloud_hass)
-
-    assert auth.cognito is None
-
-
-def test_load_auth_with_valid_auth(cloud_hass, mock_read, mock_cognito):
-    """Test calling load_auth when valid auth."""
-    auth = auth_api.load_auth(cloud_hass)
-
-    assert auth.cognito is not None
-
-
-def test_auth_properties():
-    """Test Auth class properties."""
-    auth = auth_api.Auth(None, None)
-    assert not auth.is_logged_in
-    auth.account = {}
-    assert auth.is_logged_in
-
-
-def test_auth_validate_auth_verification_fails(mock_cognito):
-    """Test validate authentication with verify request failing."""
-    mock_cognito.get_user.side_effect = aws_error('UserNotFoundException')
-
-    auth = auth_api.Auth(None, mock_cognito)
-    assert auth.validate_auth() is False
-
-
-def test_auth_validate_auth_token_refresh_needed_fails(mock_cognito):
-    """Test validate authentication with refresh needed which gets 401."""
-    mock_cognito.get_user.side_effect = aws_error('NotAuthorizedException')
-    mock_cognito.renew_access_token.side_effect = \
-        aws_error('NotAuthorizedException')
-
-    auth = auth_api.Auth(None, mock_cognito)
-    assert auth.validate_auth() is False
-
-
-def test_auth_validate_auth_token_refresh_needed_succeeds(mock_write,
-                                                          mock_cognito):
-    """Test validate authentication with refresh."""
-    mock_cognito.get_user.side_effect = [
-        aws_error('NotAuthorizedException'),
-        MagicMock(email='hello@home-assistant.io')
-    ]
-
-    auth = auth_api.Auth(None, mock_cognito)
-    assert auth.validate_auth() is True
-    assert len(mock_write.mock_calls) == 1
-
-
-def test_auth_login_invalid_auth(mock_cognito, mock_write):
     """Test trying to login with invalid credentials."""
+    cloud = MagicMock(is_logged_in=False)
     mock_cognito.authenticate.side_effect = aws_error('NotAuthorizedException')
-    auth = auth_api.Auth(None, None)
+
     with pytest.raises(auth_api.Unauthenticated):
-        auth.login('user', 'pass')
+        auth_api.login(cloud, 'user', 'pass')
 
-    assert not auth.is_logged_in
+    assert len(cloud.write_user_info.mock_calls) == 0
-    assert len(mock_cognito.get_user.mock_calls) == 0
-    assert len(mock_write.mock_calls) == 0
 
 
-def test_auth_login_user_not_found(mock_cognito, mock_write):
+def test_login_user_not_found(mock_cognito):
     """Test trying to login with invalid credentials."""
+    cloud = MagicMock(is_logged_in=False)
     mock_cognito.authenticate.side_effect = aws_error('UserNotFoundException')
-    auth = auth_api.Auth(None, None)
+
     with pytest.raises(auth_api.UserNotFound):
-        auth.login('user', 'pass')
+        auth_api.login(cloud, 'user', 'pass')
 
-    assert not auth.is_logged_in
+    assert len(cloud.write_user_info.mock_calls) == 0
-    assert len(mock_cognito.get_user.mock_calls) == 0
-    assert len(mock_write.mock_calls) == 0
 
 
-def test_auth_login_user_not_confirmed(mock_cognito, mock_write):
+def test_login_user_not_confirmed(mock_cognito):
     """Test trying to login without confirming account."""
+    cloud = MagicMock(is_logged_in=False)
     mock_cognito.authenticate.side_effect = \
         aws_error('UserNotConfirmedException')
-    auth = auth_api.Auth(None, None)
+
     with pytest.raises(auth_api.UserNotConfirmed):
-        auth.login('user', 'pass')
+        auth_api.login(cloud, 'user', 'pass')
 
-    assert not auth.is_logged_in
+    assert len(cloud.write_user_info.mock_calls) == 0
-    assert len(mock_cognito.get_user.mock_calls) == 0
-    assert len(mock_write.mock_calls) == 0
 
 
-def test_auth_login(cloud_hass, mock_cognito, mock_write):
+def test_login(mock_cognito):
     """Test trying to login without confirming account."""
-    mock_cognito.get_user.return_value = \
+    cloud = MagicMock(is_logged_in=False)
-        MagicMock(email='hello@home-assistant.io')
+    mock_cognito.id_token = 'test_id_token'
-    auth = auth_api.Auth(cloud_hass, None)
+    mock_cognito.access_token = 'test_access_token'
-    auth.login('user', 'pass')
+    mock_cognito.refresh_token = 'test_refresh_token'
-    assert auth.is_logged_in
+
+    auth_api.login(cloud, 'user', 'pass')
+
     assert len(mock_cognito.authenticate.mock_calls) == 1
-    assert len(mock_write.mock_calls) == 1
+    assert cloud.email == 'user'
-    result_hass, result_auth = mock_write.mock_calls[0][1]
+    assert cloud.id_token == 'test_id_token'
-    assert result_hass is cloud_hass
+    assert cloud.access_token == 'test_access_token'
-    assert result_auth is auth
+    assert cloud.refresh_token == 'test_refresh_token'
-
+    assert len(cloud.write_user_info.mock_calls) == 1
-
-def test_auth_renew_access_token(mock_write, mock_cognito):
-    """Test renewing an access token."""
-    auth = auth_api.Auth(None, mock_cognito)
-    assert auth.renew_access_token()
-    assert len(mock_write.mock_calls) == 1
-
-
-def test_auth_renew_access_token_fails(mock_write, mock_cognito):
-    """Test failing to renew an access token."""
-    mock_cognito.renew_access_token.side_effect = aws_error('SomeError')
-    auth = auth_api.Auth(None, mock_cognito)
-    assert not auth.renew_access_token()
-    assert len(mock_write.mock_calls) == 0
-
-
-def test_auth_logout(mock_write, mock_cognito):
-    """Test renewing an access token."""
-    auth = auth_api.Auth(None, mock_cognito)
-    auth.account = MagicMock()
-    auth.logout()
-    assert auth.account is None
-    assert len(mock_write.mock_calls) == 1
-
-
-def test_auth_logout_fails(mock_write, mock_cognito):
-    """Test error while logging out."""
-    mock_cognito.logout.side_effect = aws_error('SomeError')
-    auth = auth_api.Auth(None, mock_cognito)
-    auth.account = MagicMock()
-    with pytest.raises(auth_api.CloudError):
-        auth.logout()
-    assert auth.account is not None
-    assert len(mock_write.mock_calls) == 0
 
 
 def test_register(mock_cognito):
     """Test registering an account."""
     auth_api.register(None, 'email@home-assistant.io', 'password')
     assert len(mock_cognito.register.mock_calls) == 1
-    result_email, result_password = mock_cognito.register.mock_calls[0][1]
+    result_user, result_password = mock_cognito.register.mock_calls[0][1]
-    assert result_email == 'email@home-assistant.io'
+    assert result_user == \
+        auth_api._generate_username('email@home-assistant.io')
     assert result_password == 'password'
 
 
@@ -227,8 +97,9 @@ def test_confirm_register(mock_cognito):
     """Test confirming a registration of an account."""
     auth_api.confirm_register(None, '123456', 'email@home-assistant.io')
     assert len(mock_cognito.confirm_sign_up.mock_calls) == 1
-    result_code, result_email = mock_cognito.confirm_sign_up.mock_calls[0][1]
+    result_code, result_user = mock_cognito.confirm_sign_up.mock_calls[0][1]
-    assert result_email == 'email@home-assistant.io'
+    assert result_user == \
+        auth_api._generate_username('email@home-assistant.io')
     assert result_code == '123456'
 
 
@@ -269,3 +140,45 @@ def test_confirm_forgot_password_fails(mock_cognito):
     with pytest.raises(auth_api.CloudError):
         auth_api.confirm_forgot_password(
             None, '123456', 'email@home-assistant.io', 'new password')
+
+
+def test_check_token_writes_new_token_on_refresh(mock_cognito):
+    """Test check_token writes new token if refreshed."""
+    cloud = MagicMock()
+    mock_cognito.check_token.return_value = True
+    mock_cognito.id_token = 'new id token'
+    mock_cognito.access_token = 'new access token'
+
+    auth_api.check_token(cloud)
+
+    assert len(mock_cognito.check_token.mock_calls) == 1
+    assert cloud.id_token == 'new id token'
+    assert cloud.access_token == 'new access token'
+    assert len(cloud.write_user_info.mock_calls) == 1
+
+
+def test_check_token_does_not_write_existing_token(mock_cognito):
+    """Test check_token won't write new token if still valid."""
+    cloud = MagicMock()
+    mock_cognito.check_token.return_value = False
+
+    auth_api.check_token(cloud)
+
+    assert len(mock_cognito.check_token.mock_calls) == 1
+    assert cloud.id_token != mock_cognito.id_token
+    assert cloud.access_token != mock_cognito.access_token
+    assert len(cloud.write_user_info.mock_calls) == 0
+
+
+def test_check_token_raises(mock_cognito):
+    """Test we raise correct error."""
+    cloud = MagicMock()
+    mock_cognito.check_token.side_effect = aws_error('SomeError')
+
+    with pytest.raises(auth_api.CloudError):
+        auth_api.check_token(cloud)
+
+    assert len(mock_cognito.check_token.mock_calls) == 1
+    assert cloud.id_token != mock_cognito.id_token
+    assert cloud.access_token != mock_cognito.access_token
+    assert len(cloud.write_user_info.mock_calls) == 0

tests/components/cloud/test_http_api.py

@@ -7,25 +7,25 @@ import pytest
 from homeassistant.bootstrap import async_setup_component
 from homeassistant.components.cloud import DOMAIN, auth_api
 
+from tests.common import mock_coro
+
 
 @pytest.fixture
 def cloud_client(hass, test_client):
     """Fixture that can fetch from the cloud client."""
+    with patch('homeassistant.components.cloud.Cloud.initialize'):
         hass.loop.run_until_complete(async_setup_component(hass, 'cloud', {
             'cloud': {
-            'mode': 'development'
+                'mode': 'development',
+                'cognito_client_id': 'cognito_client_id',
+                'user_pool_id': 'user_pool_id',
+                'region': 'region',
+                'relayer': 'relayer',
             }
         }))
     return hass.loop.run_until_complete(test_client(hass.http.app))
 
 
-@pytest.fixture
-def mock_auth(cloud_client, hass):
-    """Fixture to mock authentication."""
-    auth = hass.data[DOMAIN]['auth'] = MagicMock()
-    return auth
-
-
 @pytest.fixture
 def mock_cognito():
     """Mock warrant."""
@@ -41,9 +41,9 @@ def test_account_view_no_account(cloud_client):
 
 
 @asyncio.coroutine
-def test_account_view(mock_auth, cloud_client):
+def test_account_view(hass, cloud_client):
     """Test fetching account if no account available."""
-    mock_auth.account = MagicMock(email='hello@home-assistant.io')
+    hass.data[DOMAIN].email = 'hello@home-assistant.io'
     req = yield from cloud_client.get('/api/cloud/account')
     assert req.status == 200
     result = yield from req.json()
@@ -51,9 +51,13 @@ def test_account_view(mock_auth, cloud_client):
 
 
 @asyncio.coroutine
-def test_login_view(mock_auth, cloud_client):
+def test_login_view(hass, cloud_client):
     """Test logging in."""
-    mock_auth.account = MagicMock(email='hello@home-assistant.io')
+    hass.data[DOMAIN].email = 'hello@home-assistant.io'
+
+    with patch('homeassistant.components.cloud.iot.CloudIoT.connect'), \
+            patch('homeassistant.components.cloud.'
+                  'auth_api.login') as mock_login:
         req = yield from cloud_client.post('/api/cloud/login', json={
             'email': 'my_username',
             'password': 'my_password'
@@ -62,34 +66,37 @@ def test_login_view(mock_auth, cloud_client):
     assert req.status == 200
     result = yield from req.json()
     assert result == {'email': 'hello@home-assistant.io'}
-    assert len(mock_auth.login.mock_calls) == 1
+    assert len(mock_login.mock_calls) == 1
-    result_user, result_pass = mock_auth.login.mock_calls[0][1]
+    cloud, result_user, result_pass = mock_login.mock_calls[0][1]
     assert result_user == 'my_username'
     assert result_pass == 'my_password'
 
 
 @asyncio.coroutine
-def test_login_view_invalid_json(mock_auth, cloud_client):
+def test_login_view_invalid_json(cloud_client):
     """Try logging in with invalid JSON."""
+    with patch('homeassistant.components.cloud.auth_api.login') as mock_login:
         req = yield from cloud_client.post('/api/cloud/login', data='Not JSON')
     assert req.status == 400
-    assert len(mock_auth.mock_calls) == 0
+    assert len(mock_login.mock_calls) == 0
 
 
 @asyncio.coroutine
-def test_login_view_invalid_schema(mock_auth, cloud_client):
+def test_login_view_invalid_schema(cloud_client):
     """Try logging in with invalid schema."""
+    with patch('homeassistant.components.cloud.auth_api.login') as mock_login:
         req = yield from cloud_client.post('/api/cloud/login', json={
             'invalid': 'schema'
         })
     assert req.status == 400
-    assert len(mock_auth.mock_calls) == 0
+    assert len(mock_login.mock_calls) == 0
 
 
 @asyncio.coroutine
-def test_login_view_request_timeout(mock_auth, cloud_client):
+def test_login_view_request_timeout(cloud_client):
     """Test request timeout while trying to log in."""
-    mock_auth.login.side_effect = asyncio.TimeoutError
+    with patch('homeassistant.components.cloud.auth_api.login',
+               side_effect=asyncio.TimeoutError):
         req = yield from cloud_client.post('/api/cloud/login', json={
             'email': 'my_username',
             'password': 'my_password'
@@ -99,9 +106,10 @@ def test_login_view_request_timeout(mock_auth, cloud_client):
 
 
 @asyncio.coroutine
-def test_login_view_invalid_credentials(mock_auth, cloud_client):
+def test_login_view_invalid_credentials(cloud_client):
     """Test logging in with invalid credentials."""
-    mock_auth.login.side_effect = auth_api.Unauthenticated
+    with patch('homeassistant.components.cloud.auth_api.login',
+               side_effect=auth_api.Unauthenticated):
         req = yield from cloud_client.post('/api/cloud/login', json={
             'email': 'my_username',
             'password': 'my_password'
@@ -111,9 +119,10 @@ def test_login_view_invalid_credentials(mock_auth, cloud_client):
 
 
 @asyncio.coroutine
-def test_login_view_unknown_error(mock_auth, cloud_client):
+def test_login_view_unknown_error(cloud_client):
     """Test unknown error while logging in."""
-    mock_auth.login.side_effect = auth_api.UnknownError
+    with patch('homeassistant.components.cloud.auth_api.login',
+               side_effect=auth_api.UnknownError):
         req = yield from cloud_client.post('/api/cloud/login', json={
             'email': 'my_username',
             'password': 'my_password'
@@ -123,27 +132,31 @@ def test_login_view_unknown_error(mock_auth, cloud_client):
 
 
 @asyncio.coroutine
-def test_logout_view(mock_auth, cloud_client):
+def test_logout_view(hass, cloud_client):
     """Test logging out."""
+    cloud = hass.data['cloud'] = MagicMock()
+    cloud.logout.return_value = mock_coro()
     req = yield from cloud_client.post('/api/cloud/logout')
     assert req.status == 200
     data = yield from req.json()
     assert data == {'message': 'ok'}
-    assert len(mock_auth.logout.mock_calls) == 1
+    assert len(cloud.logout.mock_calls) == 1
 
 
 @asyncio.coroutine
-def test_logout_view_request_timeout(mock_auth, cloud_client):
+def test_logout_view_request_timeout(hass, cloud_client):
     """Test timeout while logging out."""
-    mock_auth.logout.side_effect = asyncio.TimeoutError
+    cloud = hass.data['cloud'] = MagicMock()
+    cloud.logout.side_effect = asyncio.TimeoutError
     req = yield from cloud_client.post('/api/cloud/logout')
     assert req.status == 502
 
 
 @asyncio.coroutine
-def test_logout_view_unknown_error(mock_auth, cloud_client):
+def test_logout_view_unknown_error(hass, cloud_client):
     """Test unknown error while logging out."""
-    mock_auth.logout.side_effect = auth_api.UnknownError
+    cloud = hass.data['cloud'] = MagicMock()
+    cloud.logout.side_effect = auth_api.UnknownError
     req = yield from cloud_client.post('/api/cloud/logout')
     assert req.status == 502
 
@@ -158,7 +171,7 @@ def test_register_view(mock_cognito, cloud_client):
     assert req.status == 200
     assert len(mock_cognito.register.mock_calls) == 1
     result_email, result_pass = mock_cognito.register.mock_calls[0][1]
-    assert result_email == 'hello@bla.com'
+    assert result_email == auth_api._generate_username('hello@bla.com')
     assert result_pass == 'falcon42'
 
 
@@ -205,7 +218,7 @@ def test_confirm_register_view(mock_cognito, cloud_client):
     assert req.status == 200
     assert len(mock_cognito.confirm_sign_up.mock_calls) == 1
     result_code, result_email = mock_cognito.confirm_sign_up.mock_calls[0][1]
-    assert result_email == 'hello@bla.com'
+    assert result_email == auth_api._generate_username('hello@bla.com')
     assert result_code == '123456'
 
 

tests/components/cloud/test_init.py

@@ -0,0 +1,135 @@
+"""Test the cloud component."""
+import asyncio
+import json
+from unittest.mock import patch, MagicMock, mock_open
+
+import pytest
+
+from homeassistant.components import cloud
+
+from tests.common import mock_coro
+
+
+@pytest.fixture
+def mock_os():
+    """Mock os module."""
+    with patch('homeassistant.components.cloud.os') as os:
+        os.path.isdir.return_value = True
+        yield os
+
+
+@asyncio.coroutine
+def test_constructor_loads_info_from_constant():
+    """Test non-dev mode loads info from SERVERS constant."""
+    hass = MagicMock(data={})
+    with patch.dict(cloud.SERVERS, {
+        'beer': {
+            'cognito_client_id': 'test-cognito_client_id',
+            'user_pool_id': 'test-user_pool_id',
+            'region': 'test-region',
+            'relayer': 'test-relayer',
+        }
+    }):
+        result = yield from cloud.async_setup(hass, {
+            'cloud': {cloud.CONF_MODE: 'beer'}
+        })
+        assert result
+
+    cl = hass.data['cloud']
+    assert cl.mode == 'beer'
+    assert cl.cognito_client_id == 'test-cognito_client_id'
+    assert cl.user_pool_id == 'test-user_pool_id'
+    assert cl.region == 'test-region'
+    assert cl.relayer == 'test-relayer'
+
+
+@asyncio.coroutine
+def test_constructor_loads_info_from_config():
+    """Test non-dev mode loads info from SERVERS constant."""
+    hass = MagicMock(data={})
+
+    result = yield from cloud.async_setup(hass, {
+        'cloud': {
+            cloud.CONF_MODE: cloud.MODE_DEV,
+            'cognito_client_id': 'test-cognito_client_id',
+            'user_pool_id': 'test-user_pool_id',
+            'region': 'test-region',
+            'relayer': 'test-relayer',
+        }
+    })
+    assert result
+
+    cl = hass.data['cloud']
+    assert cl.mode == cloud.MODE_DEV
+    assert cl.cognito_client_id == 'test-cognito_client_id'
+    assert cl.user_pool_id == 'test-user_pool_id'
+    assert cl.region == 'test-region'
+    assert cl.relayer == 'test-relayer'
+
+
+@asyncio.coroutine
+def test_initialize_loads_info(mock_os, hass):
+    """Test initialize will load info from config file."""
+    mock_os.path.isfile.return_value = True
+    mopen = mock_open(read_data=json.dumps({
+        'email': 'test-email',
+        'id_token': 'test-id-token',
+        'access_token': 'test-access-token',
+        'refresh_token': 'test-refresh-token',
+    }))
+
+    cl = cloud.Cloud(hass, cloud.MODE_DEV)
+    cl.iot = MagicMock()
+    cl.iot.connect.return_value = mock_coro()
+
+    with patch('homeassistant.components.cloud.open', mopen, create=True):
+        yield from cl.initialize()
+
+    assert cl.email == 'test-email'
+    assert cl.id_token == 'test-id-token'
+    assert cl.access_token == 'test-access-token'
+    assert cl.refresh_token == 'test-refresh-token'
+    assert len(cl.iot.connect.mock_calls) == 1
+
+
+@asyncio.coroutine
+def test_logout_clears_info(mock_os, hass):
+    """Test logging out disconnects and removes info."""
+    cl = cloud.Cloud(hass, cloud.MODE_DEV)
+    cl.iot = MagicMock()
+    cl.iot.disconnect.return_value = mock_coro()
+
+    yield from cl.logout()
+
+    assert len(cl.iot.disconnect.mock_calls) == 1
+    assert cl.email is None
+    assert cl.id_token is None
+    assert cl.access_token is None
+    assert cl.refresh_token is None
+    assert len(mock_os.remove.mock_calls) == 1
+
+
+@asyncio.coroutine
+def test_write_user_info():
+    """Test writing user info works."""
+    mopen = mock_open()
+
+    cl = cloud.Cloud(MagicMock(), cloud.MODE_DEV)
+    cl.email = 'test-email'
+    cl.id_token = 'test-id-token'
+    cl.access_token = 'test-access-token'
+    cl.refresh_token = 'test-refresh-token'
+
+    with patch('homeassistant.components.cloud.open', mopen, create=True):
+        cl.write_user_info()
+
+    handle = mopen()
+
+    assert len(handle.write.mock_calls) == 1
+    data = json.loads(handle.write.mock_calls[0][1][0])
+    assert data == {
+        'access_token': 'test-access-token',
+        'email': 'test-email',
+        'id_token': 'test-id-token',
+        'refresh_token': 'test-refresh-token',
+    }

tests/components/cloud/test_iot.py

@@ -0,0 +1,243 @@
+"""Test the cloud.iot module."""
+import asyncio
+from unittest.mock import patch, MagicMock, PropertyMock
+
+from aiohttp import WSMsgType, client_exceptions
+import pytest
+
+from homeassistant.components.cloud import iot, auth_api
+from tests.common import mock_coro
+
+
+@pytest.fixture
+def mock_client():
+    """Mock the IoT client."""
+    client = MagicMock()
+    type(client).closed = PropertyMock(side_effect=[False, True])
+
+    with patch('asyncio.sleep'), \
+            patch('homeassistant.components.cloud.iot'
+                  '.async_get_clientsession') as session:
+        session().ws_connect.return_value = mock_coro(client)
+        yield client
+
+
+@pytest.fixture
+def mock_handle_message():
+    """Mock handle message."""
+    with patch('homeassistant.components.cloud.iot'
+               '.async_handle_message') as mock:
+        yield mock
+
+
+@asyncio.coroutine
+def test_cloud_calling_handler(mock_client, mock_handle_message):
+    """Test we call handle message with correct info."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.return_value = mock_coro(MagicMock(
+        type=WSMsgType.text,
+        json=MagicMock(return_value={
+            'msgid': 'test-msg-id',
+            'handler': 'test-handler',
+            'payload': 'test-payload'
+        })
+    ))
+    mock_handle_message.return_value = mock_coro('response')
+    mock_client.send_json.return_value = mock_coro(None)
+
+    yield from conn.connect()
+
+    # Check that we sent message to handler correctly
+    assert len(mock_handle_message.mock_calls) == 1
+    p_hass, p_cloud, handler_name, payload = \
+        mock_handle_message.mock_calls[0][1]
+
+    assert p_hass is cloud.hass
+    assert p_cloud is cloud
+    assert handler_name == 'test-handler'
+    assert payload == 'test-payload'
+
+    # Check that we forwarded response from handler to cloud
+    assert len(mock_client.send_json.mock_calls) == 1
+    assert mock_client.send_json.mock_calls[0][1][0] == {
+        'msgid': 'test-msg-id',
+        'payload': 'response'
+    }
+
+
+@asyncio.coroutine
+def test_connection_msg_for_unknown_handler(mock_client):
+    """Test a msg for an unknown handler."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.return_value = mock_coro(MagicMock(
+        type=WSMsgType.text,
+        json=MagicMock(return_value={
+            'msgid': 'test-msg-id',
+            'handler': 'non-existing-handler',
+            'payload': 'test-payload'
+        })
+    ))
+    mock_client.send_json.return_value = mock_coro(None)
+
+    yield from conn.connect()
+
+    # Check that we sent the correct error
+    assert len(mock_client.send_json.mock_calls) == 1
+    assert mock_client.send_json.mock_calls[0][1][0] == {
+        'msgid': 'test-msg-id',
+        'error': 'unknown-handler',
+    }
+
+
+@asyncio.coroutine
+def test_connection_msg_for_handler_raising(mock_client, mock_handle_message):
+    """Test we sent error when handler raises exception."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.return_value = mock_coro(MagicMock(
+        type=WSMsgType.text,
+        json=MagicMock(return_value={
+            'msgid': 'test-msg-id',
+            'handler': 'test-handler',
+            'payload': 'test-payload'
+        })
+    ))
+    mock_handle_message.side_effect = Exception('Broken')
+    mock_client.send_json.return_value = mock_coro(None)
+
+    yield from conn.connect()
+
+    # Check that we sent the correct error
+    assert len(mock_client.send_json.mock_calls) == 1
+    assert mock_client.send_json.mock_calls[0][1][0] == {
+        'msgid': 'test-msg-id',
+        'error': 'exception',
+    }
+
+
+@asyncio.coroutine
+def test_handler_forwarding():
+    """Test we forward messages to correct handler."""
+    handler = MagicMock()
+    handler.return_value = mock_coro()
+    hass = object()
+    cloud = object()
+    with patch.dict(iot.HANDLERS, {'test': handler}):
+        yield from iot.async_handle_message(
+            hass, cloud, 'test', 'payload')
+
+    assert len(handler.mock_calls) == 1
+    r_hass, r_cloud, payload = handler.mock_calls[0][1]
+    assert r_hass is hass
+    assert r_cloud is cloud
+    assert payload == 'payload'
+
+
+@asyncio.coroutine
+def test_handling_core_messages(hass):
+    """Test handling core messages."""
+    cloud = MagicMock()
+    cloud.logout.return_value = mock_coro()
+    yield from iot.async_handle_cloud(hass, cloud, {
+        'action': 'logout',
+        'reason': 'Logged in at two places.'
+    })
+    assert len(cloud.logout.mock_calls) == 1
+
+
+@asyncio.coroutine
+def test_cloud_getting_disconnected_by_server(mock_client, caplog):
+    """Test server disconnecting instance."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.return_value = mock_coro(MagicMock(
+        type=WSMsgType.CLOSING,
+    ))
+
+    yield from conn.connect()
+
+    assert 'Connection closed: Closed by server' in caplog.text
+    assert 'connect' in str(cloud.hass.async_add_job.mock_calls[-1][1][0])
+
+
+@asyncio.coroutine
+def test_cloud_receiving_bytes(mock_client, caplog):
+    """Test server disconnecting instance."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.return_value = mock_coro(MagicMock(
+        type=WSMsgType.BINARY,
+    ))
+
+    yield from conn.connect()
+
+    assert 'Connection closed: Received non-Text message' in caplog.text
+    assert 'connect' in str(cloud.hass.async_add_job.mock_calls[-1][1][0])
+
+
+@asyncio.coroutine
+def test_cloud_sending_invalid_json(mock_client, caplog):
+    """Test cloud sending invalid JSON."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.return_value = mock_coro(MagicMock(
+        type=WSMsgType.TEXT,
+        json=MagicMock(side_effect=ValueError)
+    ))
+
+    yield from conn.connect()
+
+    assert 'Connection closed: Received invalid JSON.' in caplog.text
+    assert 'connect' in str(cloud.hass.async_add_job.mock_calls[-1][1][0])
+
+
+@asyncio.coroutine
+def test_cloud_check_token_raising(mock_client, caplog):
+    """Test cloud sending invalid JSON."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.side_effect = auth_api.CloudError
+
+    yield from conn.connect()
+
+    assert 'Unable to connect: Unable to refresh token.' in caplog.text
+    assert 'connect' in str(cloud.hass.async_add_job.mock_calls[-1][1][0])
+
+
+@asyncio.coroutine
+def test_cloud_connect_invalid_auth(mock_client, caplog):
+    """Test invalid auth detected by server."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.side_effect = \
+        client_exceptions.WSServerHandshakeError(None, None, code=401)
+
+    yield from conn.connect()
+
+    assert 'Connection closed: Invalid auth.' in caplog.text
+
+
+@asyncio.coroutine
+def test_cloud_unable_to_connect(mock_client, caplog):
+    """Test unable to connect error."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.side_effect = client_exceptions.ClientError(None, None)
+
+    yield from conn.connect()
+
+    assert 'Unable to connect:' in caplog.text
+
+
+@asyncio.coroutine
+def test_cloud_random_exception(mock_client, caplog):
+    """Test random exception."""
+    cloud = MagicMock()
+    conn = iot.CloudIoT(cloud)
+    mock_client.receive.side_effect = Exception
+
+    yield from conn.connect()
+
+    assert 'Unexpected error' in caplog.text