Cloud cognito switch (#10823)

* Allow email based cognito instance

* Fix quitting Home Assistant while reconnecting

* Lint

homeassistant/components/cloud/__init__.py

@@ -104,6 +104,11 @@ class Cloud:
             self.region = info['region']
             self.relayer = info['relayer']
 
+    @property
+    def cognito_email_based(self):
+        """Return if cognito is email based."""
+        return not self.user_pool_id.endswith('GmV')
+
     @property
     def is_logged_in(self):
         """Get if cloud is logged in."""

homeassistant/components/cloud/auth_api.py

@@ -69,7 +69,10 @@ def register(cloud, email, password):
 
     cognito = _cognito(cloud)
     try:
-        cognito.register(_generate_username(email), password, email=email)
+        if cloud.cognito_email_based:
+            cognito.register(email, password, email=email)
+        else:
+            cognito.register(_generate_username(email), password, email=email)
     except ClientError as err:
         raise _map_aws_exception(err)
 
@@ -80,7 +83,11 @@ def confirm_register(cloud, confirmation_code, email):
 
     cognito = _cognito(cloud)
     try:
-        cognito.confirm_sign_up(confirmation_code, _generate_username(email))
+        if cloud.cognito_email_based:
+            cognito.confirm_sign_up(confirmation_code, email)
+        else:
+            cognito.confirm_sign_up(confirmation_code,
+                                    _generate_username(email))
     except ClientError as err:
         raise _map_aws_exception(err)
 
@@ -89,7 +96,11 @@ def forgot_password(cloud, email):
     """Initiate forgotten password flow."""
     from botocore.exceptions import ClientError
 
-    cognito = _cognito(cloud, username=_generate_username(email))
+    if cloud.cognito_email_based:
+        cognito = _cognito(cloud, username=email)
+    else:
+        cognito = _cognito(cloud, username=_generate_username(email))
+
     try:
         cognito.initiate_forgot_password()
     except ClientError as err:
@@ -100,7 +111,11 @@ def confirm_forgot_password(cloud, confirmation_code, email, new_password):
     """Confirm forgotten password code and change password."""
     from botocore.exceptions import ClientError
 
-    cognito = _cognito(cloud, username=_generate_username(email))
+    if cloud.cognito_email_based:
+        cognito = _cognito(cloud, username=email)
+    else:
+        cognito = _cognito(cloud, username=_generate_username(email))
+
     try:
         cognito.confirm_forgot_password(confirmation_code, new_password)
     except ClientError as err:

homeassistant/components/cloud/iot.py

@@ -59,13 +59,6 @@ class CloudIoT:
         if self.state == STATE_CONNECTED:
             raise RuntimeError('Already connected')
 
-        self.state = STATE_CONNECTING
-        self.close_requested = False
-        remove_hass_stop_listener = None
-        session = async_get_clientsession(self.cloud.hass)
-        client = None
-        disconnect_warn = None
-
         @asyncio.coroutine
         def _handle_hass_stop(event):
             """Handle Home Assistant shutting down."""
@@ -73,6 +66,14 @@ class CloudIoT:
             remove_hass_stop_listener = None
             yield from self.disconnect()
 
+        self.state = STATE_CONNECTING
+        self.close_requested = False
+        remove_hass_stop_listener = hass.bus.async_listen_once(
+            EVENT_HOMEASSISTANT_STOP, _handle_hass_stop)
+        session = async_get_clientsession(self.cloud.hass)
+        client = None
+        disconnect_warn = None
+
         try:
             yield from hass.async_add_job(auth_api.check_token, self.cloud)
 
@@ -83,9 +84,6 @@ class CloudIoT:
                 })
             self.tries = 0
 
-            remove_hass_stop_listener = hass.bus.async_listen_once(
-                EVENT_HOMEASSISTANT_STOP, _handle_hass_stop)
-
             _LOGGER.info('Connected')
             self.state = STATE_CONNECTED
 

tests/components/cloud/test_auth_api.py

@@ -77,7 +77,11 @@ def test_login(mock_cognito):
 
 def test_register(mock_cognito):
     """Test registering an account."""
-    auth_api.register(None, 'email@home-assistant.io', 'password')
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
+    auth_api.register(cloud, 'email@home-assistant.io', 'password')
     assert len(mock_cognito.register.mock_calls) == 1
     result_user, result_password = mock_cognito.register.mock_calls[0][1]
     assert result_user == \
@@ -87,14 +91,18 @@ def test_register(mock_cognito):
 
 def test_register_fails(mock_cognito):
     """Test registering an account."""
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
     mock_cognito.register.side_effect = aws_error('SomeError')
     with pytest.raises(auth_api.CloudError):
-        auth_api.register(None, 'email@home-assistant.io', 'password')
+        auth_api.register(cloud, 'email@home-assistant.io', 'password')
 
 
 def test_confirm_register(mock_cognito):
     """Test confirming a registration of an account."""
-    auth_api.confirm_register(None, '123456', 'email@home-assistant.io')
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
+    auth_api.confirm_register(cloud, '123456', 'email@home-assistant.io')
     assert len(mock_cognito.confirm_sign_up.mock_calls) == 1
     result_code, result_user = mock_cognito.confirm_sign_up.mock_calls[0][1]
     assert result_user == \
@@ -104,28 +112,36 @@ def test_confirm_register(mock_cognito):
 
 def test_confirm_register_fails(mock_cognito):
     """Test an error during confirmation of an account."""
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
     mock_cognito.confirm_sign_up.side_effect = aws_error('SomeError')
     with pytest.raises(auth_api.CloudError):
-        auth_api.confirm_register(None, '123456', 'email@home-assistant.io')
+        auth_api.confirm_register(cloud, '123456', 'email@home-assistant.io')
 
 
 def test_forgot_password(mock_cognito):
     """Test starting forgot password flow."""
-    auth_api.forgot_password(None, 'email@home-assistant.io')
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
+    auth_api.forgot_password(cloud, 'email@home-assistant.io')
     assert len(mock_cognito.initiate_forgot_password.mock_calls) == 1
 
 
 def test_forgot_password_fails(mock_cognito):
     """Test failure when starting forgot password flow."""
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
     mock_cognito.initiate_forgot_password.side_effect = aws_error('SomeError')
     with pytest.raises(auth_api.CloudError):
-        auth_api.forgot_password(None, 'email@home-assistant.io')
+        auth_api.forgot_password(cloud, 'email@home-assistant.io')
 
 
 def test_confirm_forgot_password(mock_cognito):
     """Test confirming forgot password."""
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
     auth_api.confirm_forgot_password(
-        None, '123456', 'email@home-assistant.io', 'new password')
+        cloud, '123456', 'email@home-assistant.io', 'new password')
     assert len(mock_cognito.confirm_forgot_password.mock_calls) == 1
     result_code, result_password = \
         mock_cognito.confirm_forgot_password.mock_calls[0][1]
@@ -135,10 +151,12 @@ def test_confirm_forgot_password(mock_cognito):
 
 def test_confirm_forgot_password_fails(mock_cognito):
     """Test failure when confirming forgot password."""
+    cloud = MagicMock()
+    cloud.cognito_email_based = False
     mock_cognito.confirm_forgot_password.side_effect = aws_error('SomeError')
     with pytest.raises(auth_api.CloudError):
         auth_api.confirm_forgot_password(
-            None, '123456', 'email@home-assistant.io', 'new password')
+            cloud, '123456', 'email@home-assistant.io', 'new password')
 
 
 def test_check_token_writes_new_token_on_refresh(mock_cognito):

tests/components/cloud/test_http_api.py

@@ -191,7 +191,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 == auth_api._generate_username('hello@bla.com')
+    assert result_email == 'hello@bla.com'
     assert result_pass == 'falcon42'
 
 
@@ -238,7 +238,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 == auth_api._generate_username('hello@bla.com')
+    assert result_email == 'hello@bla.com'
     assert result_code == '123456'