Allow tradfri groups for new imported entries (#17310)

* Clean up leftover config schema option

* Allow import groups via new config yaml setup

* Fix and add test

* Add a test without groups for legacy import

* Change default import groups to False

* Fix I/O in test

homeassistant/components/tradfri/__init__.py

@@ -25,11 +25,11 @@ CONFIG_FILE = '.tradfri_psk.conf'
 KEY_GATEWAY = 'tradfri_gateway'
 KEY_API = 'tradfri_api'
 CONF_ALLOW_TRADFRI_GROUPS = 'allow_tradfri_groups'
-DEFAULT_ALLOW_TRADFRI_GROUPS = True
+DEFAULT_ALLOW_TRADFRI_GROUPS = False
 
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
-        vol.Inclusive(CONF_HOST, 'gateway'): cv.string,
+        vol.Optional(CONF_HOST): cv.string,
         vol.Optional(CONF_ALLOW_TRADFRI_GROUPS,
                      default=DEFAULT_ALLOW_TRADFRI_GROUPS): cv.boolean,
     })
@@ -64,13 +64,14 @@ async def async_setup(hass, config):
         ))
 
     host = conf.get(CONF_HOST)
+    import_groups = conf[CONF_ALLOW_TRADFRI_GROUPS]
 
     if host is None or host in configured_hosts or host in legacy_hosts:
         return True
 
     hass.async_create_task(hass.config_entries.flow.async_init(
         DOMAIN, context={'source': config_entries.SOURCE_IMPORT},
-        data={'host': host}
+        data={CONF_HOST: host, CONF_IMPORT_GROUPS: import_groups}
     ))
 
     return True

homeassistant/components/tradfri/config_flow.py

@@ -34,6 +34,7 @@ class FlowHandler(config_entries.ConfigFlow):
     def __init__(self):
         """Initialize flow."""
         self._host = None
+        self._import_groups = False
 
     async def async_step_user(self, user_input=None):
         """Handle a flow initialized by the user."""
@@ -52,7 +53,8 @@ class FlowHandler(config_entries.ConfigFlow):
 
                 # We don't ask for import group anymore as group state
                 # is not reliable, don't want to show that to the user.
-                auth[CONF_IMPORT_GROUPS] = False
+                # But we still allow specifying import group via config yaml.
+                auth[CONF_IMPORT_GROUPS] = self._import_groups
 
                 return await self._entry_from_data(auth)
 
@@ -97,6 +99,7 @@ class FlowHandler(config_entries.ConfigFlow):
         # Happens if user has host directly in configuration.yaml
         if 'key' not in user_input:
             self._host = user_input['host']
+            self._import_groups = user_input[CONF_IMPORT_GROUPS]
             return await self.async_step_auth()
 
         try:

tests/components/tradfri/conftest.py

@@ -0,0 +1,12 @@
+"""Common tradfri test fixtures."""
+from unittest.mock import patch
+
+import pytest
+
+
+@pytest.fixture
+def mock_gateway_info():
+    """Mock get_gateway_info."""
+    with patch('homeassistant.components.tradfri.config_flow.'
+               'get_gateway_info') as mock_gateway:
+        yield mock_gateway

tests/components/tradfri/test_config_flow.py

@@ -17,14 +17,6 @@ def mock_auth():
         yield mock_auth
 
 
-@pytest.fixture
-def mock_gateway_info():
-    """Mock get_gateway_info."""
-    with patch('homeassistant.components.tradfri.config_flow.'
-               'get_gateway_info') as mock_gateway:
-        yield mock_gateway
-
-
 @pytest.fixture
 def mock_entry_setup():
     """Mock entry setup."""
@@ -125,34 +117,65 @@ async def test_discovery_connection(hass, mock_auth, mock_entry_setup):
     }
 
 
-async def test_import_connection(hass, mock_gateway_info, mock_entry_setup):
+async def test_import_connection(hass, mock_auth, mock_entry_setup):
     """Test a connection via import."""
-    mock_gateway_info.side_effect = \
+    mock_auth.side_effect = lambda hass, host, code: mock_coro({
-        lambda hass, host, identity, key: mock_coro({
+        'host': host,
-            'host': host,
+        'gateway_id': 'bla',
-            'identity': identity,
+        'identity': 'mock-iden',
-            'key': key,
+        'key': 'mock-key',
-            'gateway_id': 'mock-gateway'
+    })
-        })
 
-    result = await hass.config_entries.flow.async_init(
+    flow = await hass.config_entries.flow.async_init(
         'tradfri', context={'source': 'import'}, data={
             'host': '123.123.123.123',
-            'identity': 'mock-iden',
-            'key': 'mock-key',
             'import_groups': True
         })
 
+    result = await hass.config_entries.flow.async_configure(flow['flow_id'], {
+        'security_code': 'abcd',
+    })
+
     assert result['type'] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
     assert result['result'].data == {
         'host': '123.123.123.123',
-        'gateway_id': 'mock-gateway',
+        'gateway_id': 'bla',
         'identity': 'mock-iden',
         'key': 'mock-key',
         'import_groups': True
     }
 
-    assert len(mock_gateway_info.mock_calls) == 1
+    assert len(mock_entry_setup.mock_calls) == 1
+
+
+async def test_import_connection_no_groups(hass, mock_auth, mock_entry_setup):
+    """Test a connection via import and no groups allowed."""
+    mock_auth.side_effect = lambda hass, host, code: mock_coro({
+        'host': host,
+        'gateway_id': 'bla',
+        'identity': 'mock-iden',
+        'key': 'mock-key',
+    })
+
+    flow = await hass.config_entries.flow.async_init(
+        'tradfri', context={'source': 'import'}, data={
+            'host': '123.123.123.123',
+            'import_groups': False
+        })
+
+    result = await hass.config_entries.flow.async_configure(flow['flow_id'], {
+        'security_code': 'abcd',
+    })
+
+    assert result['type'] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert result['result'].data == {
+        'host': '123.123.123.123',
+        'gateway_id': 'bla',
+        'identity': 'mock-iden',
+        'key': 'mock-key',
+        'import_groups': False
+    }
+
     assert len(mock_entry_setup.mock_calls) == 1
 
 
@@ -187,6 +210,37 @@ async def test_import_connection_legacy(hass, mock_gateway_info,
     assert len(mock_entry_setup.mock_calls) == 1
 
 
+async def test_import_connection_legacy_no_groups(
+        hass, mock_gateway_info, mock_entry_setup):
+    """Test a connection via legacy import and no groups allowed."""
+    mock_gateway_info.side_effect = \
+        lambda hass, host, identity, key: mock_coro({
+            'host': host,
+            'identity': identity,
+            'key': key,
+            'gateway_id': 'mock-gateway'
+        })
+
+    result = await hass.config_entries.flow.async_init(
+        'tradfri', context={'source': 'import'}, data={
+            'host': '123.123.123.123',
+            'key': 'mock-key',
+            'import_groups': False
+        })
+
+    assert result['type'] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert result['result'].data == {
+        'host': '123.123.123.123',
+        'gateway_id': 'mock-gateway',
+        'identity': 'homeassistant',
+        'key': 'mock-key',
+        'import_groups': False
+    }
+
+    assert len(mock_gateway_info.mock_calls) == 1
+    assert len(mock_entry_setup.mock_calls) == 1
+
+
 async def test_discovery_duplicate_aborted(hass):
     """Test a duplicate discovery host is ignored."""
     MockConfigEntry(

tests/components/tradfri/test_init.py

@@ -58,7 +58,7 @@ async def test_config_json_host_not_imported(hass):
     assert len(mock_init.mock_calls) == 0
 
 
-async def test_config_json_host_imported(hass):
+async def test_config_json_host_imported(hass, mock_gateway_info):
     """Test that we import a configured host."""
     with patch('homeassistant.components.tradfri.load_json',
                return_value={'mock-host': {'key': 'some-info'}}):