Improve tests based on Martins feedback from Axis improve tests (#30442)

tests/components/unifi/test_controller.py

@@ -1,17 +1,16 @@
 """Test UniFi Controller."""
 from collections import deque
+from copy import deepcopy
 from datetime import timedelta
 
 import aiounifi
 from asynctest import Mock, patch
 import pytest
 
-from homeassistant import config_entries
 from homeassistant.components import unifi
 from homeassistant.components.unifi.const import (
     CONF_CONTROLLER,
     CONF_SITE_ID,
-    UNIFI_CONFIG,
     UNIFI_WIRELESS_CLIENTS,
 )
 from homeassistant.const import (
@@ -21,7 +20,9 @@ from homeassistant.const import (
     CONF_USERNAME,
     CONF_VERIFY_SSL,
 )
-from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.setup import async_setup_component
+
+from tests.common import MockConfigEntry
 
 CONTROLLER_HOST = {
     "hostname": "controller_host",
@@ -47,35 +48,38 @@ CONTROLLER_DATA = {
 }
 
 ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA}
+ENTRY_OPTIONS = {}
+
+CONFIGURATION = []
 
 SITES = {"Site name": {"desc": "Site name", "name": "site_id", "role": "admin"}}
 
 
 async def setup_unifi_integration(
     hass,
-    config,
+    config=ENTRY_CONFIG,
-    options,
+    options=ENTRY_OPTIONS,
-    sites,
+    sites=SITES,
-    clients_response,
+    clients_response=None,
-    devices_response,
+    devices_response=None,
-    clients_all_response,
+    clients_all_response=None,
     known_wireless_clients=None,
+    controllers=None,
 ):
     """Create the UniFi controller."""
-    if UNIFI_CONFIG not in hass.data:
+    configuration = {}
-        hass.data[UNIFI_CONFIG] = []
+    if controllers:
-    hass.data[UNIFI_WIRELESS_CLIENTS] = unifi.UnifiWirelessClients(hass)
+        configuration = {unifi.DOMAIN: {unifi.CONF_CONTROLLERS: controllers}}
-    config_entry = config_entries.ConfigEntry(
+
-        version=1,
+    assert await async_setup_component(hass, unifi.DOMAIN, configuration)
+
+    config_entry = MockConfigEntry(
         domain=unifi.DOMAIN,
-        title="Mock Title",
+        data=deepcopy(config),
-        data=config,
+        options=deepcopy(options),
-        source="test",
-        connection_class=config_entries.CONN_CLASS_LOCAL_POLL,
-        system_options={},
-        options=options,
         entry_id=1,
     )
+    config_entry.add_to_hass(hass)
 
     if known_wireless_clients:
         hass.data[UNIFI_WIRELESS_CLIENTS].update_data(
@@ -83,13 +87,16 @@ async def setup_unifi_integration(
         )
 
     mock_client_responses = deque()
-    mock_client_responses.append(clients_response)
+    if clients_response:
+        mock_client_responses.append(clients_response)
 
     mock_device_responses = deque()
-    mock_device_responses.append(devices_response)
+    if devices_response:
+        mock_device_responses.append(devices_response)
 
     mock_client_all_responses = deque()
-    mock_client_all_responses.append(clients_all_response)
+    if clients_all_response:
+        mock_client_all_responses.append(clients_all_response)
 
     mock_requests = []
 
@@ -107,9 +114,8 @@ async def setup_unifi_integration(
     with patch("aiounifi.Controller.login", return_value=True), patch(
         "aiounifi.Controller.sites", return_value=sites
     ), patch("aiounifi.Controller.request", new=mock_request):
-        await unifi.async_setup_entry(hass, config_entry)
+        await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
-    hass.config_entries._entries.append(config_entry)
 
     controller_id = unifi.get_controller_id_from_config_entry(config_entry)
     if controller_id not in hass.data[unifi.DOMAIN]:
@@ -130,23 +136,15 @@ async def test_controller_setup(hass):
         "homeassistant.config_entries.ConfigEntries.async_forward_entry_setup",
         return_value=True,
     ) as forward_entry_setup:
-        controller = await setup_unifi_integration(
+        controller = await setup_unifi_integration(hass)
-            hass,
-            ENTRY_CONFIG,
-            options={},
-            sites=SITES,
-            clients_response=[],
-            devices_response=[],
-            clients_all_response=[],
-        )
 
-        entry = controller.config_entry
+    entry = controller.config_entry
-        assert len(forward_entry_setup.mock_calls) == len(
+    assert len(forward_entry_setup.mock_calls) == len(
-            unifi.controller.SUPPORTED_PLATFORMS
+        unifi.controller.SUPPORTED_PLATFORMS
-        )
+    )
-        assert forward_entry_setup.mock_calls[0][1] == (entry, "device_tracker")
+    assert forward_entry_setup.mock_calls[0][1] == (entry, "device_tracker")
-        assert forward_entry_setup.mock_calls[1][1] == (entry, "sensor")
+    assert forward_entry_setup.mock_calls[1][1] == (entry, "sensor")
-        assert forward_entry_setup.mock_calls[2][1] == (entry, "switch")
+    assert forward_entry_setup.mock_calls[2][1] == (entry, "switch")
 
     assert controller.host == CONTROLLER_DATA[CONF_HOST]
     assert controller.site == CONTROLLER_DATA[CONF_SITE_ID]
@@ -176,25 +174,16 @@ async def test_controller_setup(hass):
 
 async def test_controller_mac(hass):
     """Test that it is possible to identify controller mac."""
-    controller = await setup_unifi_integration(
+    controller = await setup_unifi_integration(hass, clients_response=[CONTROLLER_HOST])
-        hass,
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=[CONTROLLER_HOST],
-        devices_response=[],
-        clients_all_response=[],
-    )
     assert controller.mac == "10:00:00:00:00:01"
 
 
 async def test_controller_import_config(hass):
     """Test that import configuration.yaml instructions work."""
-    hass.data[UNIFI_CONFIG] = [
+    controllers = [
         {
             CONF_HOST: "1.2.3.4",
             CONF_SITE_ID: "Site name",
-            unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True,
             unifi.CONF_BLOCK_CLIENT: ["random mac"],
             unifi.CONF_DONT_TRACK_CLIENTS: True,
             unifi.CONF_DONT_TRACK_DEVICES: True,
@@ -203,15 +192,8 @@ async def test_controller_import_config(hass):
             unifi.CONF_SSID_FILTER: ["SSID"],
         }
     ]
-    controller = await setup_unifi_integration(
+
-        hass,
+    controller = await setup_unifi_integration(hass, controllers=controllers)
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
-    )
 
     assert controller.option_allow_bandwidth_sensors is False
     assert controller.option_block_clients == ["random mac"]
@@ -226,44 +208,21 @@ async def test_controller_not_accessible(hass):
     """Retry to login gets scheduled when connection fails."""
     with patch.object(
         unifi.controller, "get_controller", side_effect=unifi.errors.CannotConnect
-    ), pytest.raises(ConfigEntryNotReady):
+    ):
-        await setup_unifi_integration(
+        await setup_unifi_integration(hass)
-            hass,
+    assert hass.data[unifi.DOMAIN] == {}
-            ENTRY_CONFIG,
-            options={},
-            sites=SITES,
-            clients_response=[],
-            devices_response=[],
-            clients_all_response=[],
-        )
 
 
 async def test_controller_unknown_error(hass):
     """Unknown errors are handled."""
     with patch.object(unifi.controller, "get_controller", side_effect=Exception):
-        await setup_unifi_integration(
+        await setup_unifi_integration(hass)
-            hass,
+    assert hass.data[unifi.DOMAIN] == {}
-            ENTRY_CONFIG,
-            options={},
-            sites=SITES,
-            clients_response=[],
-            devices_response=[],
-            clients_all_response=[],
-        )
-        assert hass.data[unifi.DOMAIN] == {}
 
 
 async def test_reset_after_successful_setup(hass):
     """Calling reset when the entry has been setup."""
-    controller = await setup_unifi_integration(
+    controller = await setup_unifi_integration(hass)
-        hass,
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
-    )
 
     assert len(controller.listeners) == 5
 
@@ -276,15 +235,7 @@ async def test_reset_after_successful_setup(hass):
 
 async def test_failed_update_failed_login(hass):
     """Running update can handle a failed login."""
-    controller = await setup_unifi_integration(
+    controller = await setup_unifi_integration(hass)
-        hass,
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
-    )
 
     with patch.object(
         controller.api.clients, "update", side_effect=aiounifi.LoginRequired
@@ -297,15 +248,7 @@ async def test_failed_update_failed_login(hass):
 
 async def test_failed_update_successful_login(hass):
     """Running update can login when requested."""
-    controller = await setup_unifi_integration(
+    controller = await setup_unifi_integration(hass)
-        hass,
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
-    )
 
     with patch.object(
         controller.api.clients, "update", side_effect=aiounifi.LoginRequired
@@ -318,15 +261,7 @@ async def test_failed_update_successful_login(hass):
 
 async def test_failed_update(hass):
     """Running update can login when requested."""
-    controller = await setup_unifi_integration(
+    controller = await setup_unifi_integration(hass)
-        hass,
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
-    )
 
     with patch.object(
         controller.api.clients, "update", side_effect=aiounifi.AiounifiException
@@ -357,32 +292,23 @@ async def test_get_controller_verify_ssl_false(hass):
 
 async def test_get_controller_login_failed(hass):
     """Check that get_controller can handle a failed login."""
-    result = None
+    with patch(
-    with patch("aiounifi.Controller.login", side_effect=aiounifi.Unauthorized):
+        "aiounifi.Controller.login", side_effect=aiounifi.Unauthorized
-        try:
+    ), pytest.raises(unifi.errors.AuthenticationRequired):
-            result = await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
+        await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
-        except unifi.errors.AuthenticationRequired:
-            pass
-        assert result is None
 
 
 async def test_get_controller_controller_unavailable(hass):
     """Check that get_controller can handle controller being unavailable."""
-    result = None
+    with patch(
-    with patch("aiounifi.Controller.login", side_effect=aiounifi.RequestError):
+        "aiounifi.Controller.login", side_effect=aiounifi.RequestError
-        try:
+    ), pytest.raises(unifi.errors.CannotConnect):
-            result = await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
+        await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
-        except unifi.errors.CannotConnect:
-            pass
-        assert result is None
 
 
 async def test_get_controller_unknown_error(hass):
     """Check that get_controller can handle unkown errors."""
-    result = None
+    with patch(
-    with patch("aiounifi.Controller.login", side_effect=aiounifi.AiounifiException):
+        "aiounifi.Controller.login", side_effect=aiounifi.AiounifiException
-        try:
+    ), pytest.raises(unifi.errors.AuthenticationRequired):
-            result = await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
+        await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
-        except unifi.errors.AuthenticationRequired:
-            pass
-        assert result is None

tests/components/unifi/test_device_tracker.py

@@ -17,7 +17,7 @@ from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
-from .test_controller import ENTRY_CONFIG, SITES, setup_unifi_integration
+from .test_controller import ENTRY_CONFIG, setup_unifi_integration
 
 CLIENT_1 = {
     "essid": "ssid",
@@ -95,15 +95,7 @@ async def test_platform_manually_configured(hass):
 
 async def test_no_clients(hass):
     """Test the update_clients function when no clients are found."""
-    await setup_unifi_integration(
+    await setup_unifi_integration(hass)
-        hass,
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
-    )
 
     assert len(hass.states.async_all()) == 2
 
@@ -115,12 +107,9 @@ async def test_tracked_devices(hass):
 
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={CONF_SSID_FILTER: ["ssid"]},
-        sites=SITES,
         clients_response=[CLIENT_1, CLIENT_2, CLIENT_3, client_4_copy],
         devices_response=[DEVICE_1, DEVICE_2],
-        clients_all_response={},
         known_wireless_clients=(CLIENT_4["mac"],),
     )
     assert len(hass.states.async_all()) == 6
@@ -196,15 +185,7 @@ async def test_wireless_client_go_wired_issue(hass):
     client_1_client = copy(CLIENT_1)
     client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
 
-    controller = await setup_unifi_integration(
+    controller = await setup_unifi_integration(hass, clients_response=[client_1_client])
-        hass,
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=[client_1_client],
-        devices_response=[],
-        clients_all_response=[],
-    )
     assert len(hass.states.async_all()) == 3
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -277,11 +258,8 @@ async def test_restoring_client(hass):
 
     await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={unifi.CONF_BLOCK_CLIENT: True},
-        sites=SITES,
         clients_response=[CLIENT_2],
-        devices_response=[],
         clients_all_response=[CLIENT_1],
     )
     assert len(hass.states.async_all()) == 4
@@ -294,12 +272,9 @@ async def test_dont_track_clients(hass):
     """Test dont track clients config works."""
     await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={unifi.controller.CONF_TRACK_CLIENTS: False},
-        sites=SITES,
         clients_response=[CLIENT_1],
         devices_response=[DEVICE_1],
-        clients_all_response=[],
     )
     assert len(hass.states.async_all()) == 3
 
@@ -315,12 +290,9 @@ async def test_dont_track_devices(hass):
     """Test dont track devices config works."""
     await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={unifi.controller.CONF_TRACK_DEVICES: False},
-        sites=SITES,
         clients_response=[CLIENT_1],
         devices_response=[DEVICE_1],
-        clients_all_response=[],
     )
     assert len(hass.states.async_all()) == 3
 
@@ -336,12 +308,8 @@ async def test_dont_track_wired_clients(hass):
     """Test dont track wired clients config works."""
     await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={unifi.controller.CONF_TRACK_WIRED_CLIENTS: False},
-        sites=SITES,
         clients_response=[CLIENT_1, CLIENT_2],
-        devices_response=[],
-        clients_all_response=[],
     )
     assert len(hass.states.async_all()) == 3
 

tests/components/unifi/test_sensor.py

@@ -5,7 +5,7 @@ from homeassistant.components import unifi
 import homeassistant.components.sensor as sensor
 from homeassistant.setup import async_setup_component
 
-from .test_controller import ENTRY_CONFIG, SITES, setup_unifi_integration
+from .test_controller import setup_unifi_integration
 
 CLIENTS = [
     {
@@ -51,13 +51,7 @@ async def test_platform_manually_configured(hass):
 async def test_no_clients(hass):
     """Test the update_clients function when no clients are found."""
     controller = await setup_unifi_integration(
-        hass,
+        hass, options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
-        ENTRY_CONFIG,
-        options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
     )
 
     assert len(controller.mock_requests) == 3
@@ -68,16 +62,12 @@ async def test_sensors(hass):
     """Test the update_items function with some clients."""
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={
             unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True,
             unifi.const.CONF_TRACK_CLIENTS: False,
             unifi.const.CONF_TRACK_DEVICES: False,
         },
-        sites=SITES,
         clients_response=CLIENTS,
-        devices_response=[],
-        clients_all_response=[],
     )
 
     assert len(controller.mock_requests) == 3

tests/components/unifi/test_switch.py

@@ -201,15 +201,10 @@ async def test_no_clients(hass):
     """Test the update_clients function when no clients are found."""
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={
             unifi.const.CONF_TRACK_CLIENTS: False,
             unifi.const.CONF_TRACK_DEVICES: False,
         },
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
     )
 
     assert len(controller.mock_requests) == 3
@@ -220,15 +215,12 @@ async def test_controller_not_client(hass):
     """Test that the controller doesn't become a switch."""
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={
             unifi.const.CONF_TRACK_CLIENTS: False,
             unifi.const.CONF_TRACK_DEVICES: False,
         },
-        sites=SITES,
         clients_response=[CONTROLLER_HOST],
         devices_response=[DEVICE_1],
-        clients_all_response=[],
     )
 
     assert len(controller.mock_requests) == 3
@@ -243,7 +235,6 @@ async def test_not_admin(hass):
     sites["Site name"]["role"] = "not admin"
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={
             unifi.const.CONF_TRACK_CLIENTS: False,
             unifi.const.CONF_TRACK_DEVICES: False,
@@ -251,7 +242,6 @@ async def test_not_admin(hass):
         sites=sites,
         clients_response=[CLIENT_1],
         devices_response=[DEVICE_1],
-        clients_all_response=[],
     )
 
     assert len(controller.mock_requests) == 3
@@ -262,13 +252,11 @@ async def test_switches(hass):
     """Test the update_items function with some clients."""
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={
             unifi.CONF_BLOCK_CLIENT: [BLOCKED["mac"], UNBLOCKED["mac"]],
             unifi.const.CONF_TRACK_CLIENTS: False,
             unifi.const.CONF_TRACK_DEVICES: False,
         },
-        sites=SITES,
         clients_response=[CLIENT_1, CLIENT_4],
         devices_response=[DEVICE_1],
         clients_all_response=[BLOCKED, UNBLOCKED, CLIENT_1],
@@ -301,15 +289,11 @@ async def test_new_client_discovered_on_block_control(hass):
     """Test if 2nd update has a new client."""
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={
             unifi.CONF_BLOCK_CLIENT: [BLOCKED["mac"]],
             unifi.const.CONF_TRACK_CLIENTS: False,
             unifi.const.CONF_TRACK_DEVICES: False,
         },
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
         clients_all_response=[BLOCKED],
     )
 
@@ -345,15 +329,12 @@ async def test_new_client_discovered_on_poe_control(hass):
     """Test if 2nd update has a new client."""
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={
             unifi.const.CONF_TRACK_CLIENTS: False,
             unifi.const.CONF_TRACK_DEVICES: False,
         },
-        sites=SITES,
         clients_response=[CLIENT_1],
         devices_response=[DEVICE_1],
-        clients_all_response=[],
     )
 
     assert len(controller.mock_requests) == 3
@@ -402,13 +383,7 @@ async def test_ignore_multiple_poe_clients_on_same_port(hass):
     clients will be transparently marked as having POE as well.
     """
     controller = await setup_unifi_integration(
-        hass,
+        hass, clients_response=POE_SWITCH_CLIENTS, devices_response=[DEVICE_1],
-        ENTRY_CONFIG,
-        options={},
-        sites=SITES,
-        clients_response=POE_SWITCH_CLIENTS,
-        devices_response=[DEVICE_1],
-        clients_all_response=[],
     )
 
     assert len(controller.mock_requests) == 3
@@ -452,13 +427,11 @@ async def test_restoring_client(hass):
 
     controller = await setup_unifi_integration(
         hass,
-        ENTRY_CONFIG,
         options={
             unifi.CONF_BLOCK_CLIENT: ["random mac"],
             unifi.const.CONF_TRACK_CLIENTS: False,
             unifi.const.CONF_TRACK_DEVICES: False,
         },
-        sites=SITES,
         clients_response=[CLIENT_2],
         devices_response=[DEVICE_1],
         clients_all_response=[CLIENT_1],