UniFi - Make clients proper push based (#35273)

* Improve client tracker to be more comprehensible and streamlined
Improve block switches
Improve tests

* Small clean up

* Add descriptions on ssid test

* Improve test

* Make polling default off, only POE clients left to verify

* Minor improvements

* On removal cancel scheduled updates

* POE works without polling now

* Combine else and if to an elif

tests/components/unifi/test_device_tracker.py

@@ -29,7 +29,6 @@ import homeassistant.util.dt as dt_util
 
 from .test_controller import ENTRY_CONFIG, setup_unifi_integration
 
-from tests.async_mock import patch
 from tests.common import async_fire_time_changed
 
 CLIENT_1 = {
@@ -167,7 +166,6 @@ async def test_tracked_wireless_clients(hass):
 
     # State change signalling works without events
     client_1_copy = copy(CLIENT_1)
-    client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     controller.api.websocket._data = {
         "meta": {"message": MESSAGE_CLIENT},
         "data": [client_1_copy],
@@ -205,8 +203,6 @@ async def test_tracked_wireless_clients(hass):
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
 
-    # test wired bug
-
 
 async def test_tracked_clients(hass):
     """Test the update_items function with some clients."""
@@ -217,10 +213,9 @@ async def test_tracked_clients(hass):
         hass,
         options={CONF_SSID_FILTER: ["ssid"]},
         clients_response=[CLIENT_1, CLIENT_2, CLIENT_3, CLIENT_5, client_4_copy],
-        devices_response=[DEVICE_1, DEVICE_2],
-        known_wireless_clients=([CLIENT_4["mac"]]),
+        known_wireless_clients=(CLIENT_4["mac"],),
     )
-    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 6
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 4
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
@@ -246,9 +241,9 @@ async def test_tracked_clients(hass):
 
     # State change signalling works
     client_1_copy = copy(CLIENT_1)
-    client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
     controller.api.message_handler(event)
+    await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
@@ -357,7 +352,7 @@ async def test_controller_state_change(hass):
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
-    assert client_1.state == "not_home"
+    assert client_1.state == "home"
 
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1.state == "home"
@@ -499,29 +494,56 @@ async def test_option_track_devices(hass):
 
 
 async def test_option_ssid_filter(hass):
-    """Test the SSID filter works."""
-    controller = await setup_unifi_integration(hass, clients_response=[CLIENT_3])
-    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
+    """Test the SSID filter works.
+
+    Client 1 will travel from a supported SSID to an unsupported ssid.
+    Client 3 will be removed on change of options since it is in an unsupported SSID.
+    """
+    client_1_copy = copy(CLIENT_1)
+    client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+
+    controller = await setup_unifi_integration(
+        hass, clients_response=[client_1_copy, CLIENT_3]
+    )
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
 
     client_3 = hass.states.get("device_tracker.client_3")
     assert client_3
 
-    # Set SSID filter
+    # Setting SSID filter will remove clients outside of filter
     hass.config_entries.async_update_entry(
         controller.config_entry, options={CONF_SSID_FILTER: ["ssid"]},
     )
     await hass.async_block_till_done()
 
+    # Not affected by SSID filter
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+
+    # Removed due to SSID filter
     client_3 = hass.states.get("device_tracker.client_3")
     assert not client_3
 
+    # Roams to SSID outside of filter
+    client_1_copy = copy(CLIENT_1)
+    client_1_copy["essid"] = "other_ssid"
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
+    controller.api.message_handler(event)
+    # Data update while SSID filter is in effect shouldn't create the client
     client_3_copy = copy(CLIENT_3)
     client_3_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
-    # SSID filter active even though time stamp should mark as home
+    # SSID filter marks client as away
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+
+    # SSID still outside of filter
     client_3 = hass.states.get("device_tracker.client_3")
     assert not client_3
 
@@ -529,13 +551,37 @@ async def test_option_ssid_filter(hass):
     hass.config_entries.async_update_entry(
         controller.config_entry, options={CONF_SSID_FILTER: []},
     )
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
+    controller.api.message_handler(event)
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+
     client_3 = hass.states.get("device_tracker.client_3")
     assert client_3.state == "home"
 
+    async_fire_time_changed(hass, dt_util.utcnow() + controller.option_detection_time)
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+
+    # Client won't go away until after next update
+    client_3 = hass.states.get("device_tracker.client_3")
+    assert client_3.state == "home"
+
+    # Trigger update to get client marked as away
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [CLIENT_3]}
+    controller.api.message_handler(event)
+    async_fire_time_changed(hass, dt_util.utcnow() + controller.option_detection_time)
+    await hass.async_block_till_done()
+
+    client_3 = hass.states.get("device_tracker.client_3")
+    assert client_3.state == "not_home"
+
 
 async def test_wireless_client_go_wired_issue(hass):
     """Test the solution to catch wireless device go wired UniFi issue.
@@ -548,38 +594,49 @@ async def test_wireless_client_go_wired_issue(hass):
     controller = await setup_unifi_integration(hass, clients_response=[client_1_client])
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
 
+    # Client is wireless
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
 
+    # Trigger wired bug
     client_1_client["is_wired"] = True
-    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    # Wired bug fix keeps client marked as wireless
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
 
-    with patch.object(
-        dt_util, "utcnow", return_value=(dt_util.utcnow() + timedelta(minutes=5)),
-    ):
-        event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
-        controller.api.message_handler(event)
-        await hass.async_block_till_done()
+    # Pass time
+    async_fire_time_changed(hass, dt_util.utcnow() + controller.option_detection_time)
+    await hass.async_block_till_done()
 
-        client_1 = hass.states.get("device_tracker.client_1")
-        assert client_1.state == "not_home"
-        assert client_1.attributes["is_wired"] is False
+    # Marked as home according to the timer
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+    assert client_1.attributes["is_wired"] is False
 
-    client_1_client["is_wired"] = False
-    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
-    event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
+    # Try to mark client as connected
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    # Make sure it don't go online again until wired bug disappears
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+    assert client_1.attributes["is_wired"] is False
+
+    # Make client wireless
+    client_1_client["is_wired"] = False
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    controller.api.message_handler(event)
+    await hass.async_block_till_done()
+
+    # Client is no longer affected by wired bug and can be marked online
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
@@ -595,27 +652,49 @@ async def test_option_ignore_wired_bug(hass):
     )
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
 
+    # Client is wireless
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
 
+    # Trigger wired bug
     client_1_client["is_wired"] = True
-    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
-    event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    # Wired bug in effect
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is True
 
-    client_1_client["is_wired"] = False
-    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
-    event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
+    # pass time
+    async_fire_time_changed(hass, dt_util.utcnow() + controller.option_detection_time)
+    await hass.async_block_till_done()
+
+    # Timer marks client as away
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+    assert client_1.attributes["is_wired"] is True
+
+    # Mark client as connected again
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    # Ignoring wired bug allows client to go home again even while affected
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+    assert client_1.attributes["is_wired"] is True
+
+    # Make client wireless
+    client_1_client["is_wired"] = False
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    controller.api.message_handler(event)
+    await hass.async_block_till_done()
+
+    # Client is wireless and still connected
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
@@ -665,7 +744,7 @@ async def test_restoring_client(hass):
 
 async def test_dont_track_clients(hass):
     """Test don't track clients config works."""
-    await setup_unifi_integration(
+    controller = await setup_unifi_integration(
         hass,
         options={CONF_TRACK_CLIENTS: False},
         clients_response=[CLIENT_1],
@@ -678,12 +757,24 @@ async def test_dont_track_clients(hass):
 
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1 is not None
-    assert device_1.state == "home"
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_CLIENTS: True},
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
 
 
 async def test_dont_track_devices(hass):
     """Test don't track devices config works."""
-    await setup_unifi_integration(
+    controller = await setup_unifi_integration(
         hass,
         options={CONF_TRACK_DEVICES: False},
         clients_response=[CLIENT_1],
@@ -693,15 +784,27 @@ async def test_dont_track_devices(hass):
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
-    assert client_1.state == "not_home"
 
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1 is None
 
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_DEVICES: True},
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
 
 async def test_dont_track_wired_clients(hass):
     """Test don't track wired clients config works."""
-    await setup_unifi_integration(
+    controller = await setup_unifi_integration(
         hass,
         options={CONF_TRACK_WIRED_CLIENTS: False},
         clients_response=[CLIENT_1, CLIENT_2],
@@ -710,7 +813,19 @@ async def test_dont_track_wired_clients(hass):
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
-    assert client_1.state == "not_home"
 
-    client_2 = hass.states.get("device_tracker.client_2")
+    client_2 = hass.states.get("device_tracker.wired_client")
     assert client_2 is None
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_WIRED_CLIENTS: True},
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None

tests/components/unifi/test_switch.py

@@ -211,6 +211,30 @@ EVENT_BLOCKED_CLIENT_CONNECTED = {
     "_id": "5ea331fa30c49e00f90ddc1a",
 }
 
+EVENT_BLOCKED_CLIENT_BLOCKED = {
+    "user": BLOCKED["mac"],
+    "hostname": BLOCKED["hostname"],
+    "key": "EVT_WC_Blocked",
+    "subsystem": "wlan",
+    "site_id": "name",
+    "time": 1587753456179,
+    "datetime": "2020-04-24T18:37:36Z",
+    "msg": f'User{[BLOCKED["mac"]]} has been blocked."',
+    "_id": "5ea331fa30c49e00f90ddc1a",
+}
+
+EVENT_BLOCKED_CLIENT_UNBLOCKED = {
+    "user": BLOCKED["mac"],
+    "hostname": BLOCKED["hostname"],
+    "key": "EVT_WC_Unblocked",
+    "subsystem": "wlan",
+    "site_id": "name",
+    "time": 1587753456179,
+    "datetime": "2020-04-24T18:37:36Z",
+    "msg": f'User{[BLOCKED["mac"]]} has been unblocked."',
+    "_id": "5ea331fa30c49e00f90ddc1a",
+}
+
 
 EVENT_CLIENT_2_CONNECTED = {
     "user": CLIENT_2["mac"],
@@ -368,6 +392,74 @@ async def test_remove_switches(hass):
     assert block_switch is None
 
 
+async def test_block_switches(hass):
+    """Test the update_items function with some clients."""
+    controller = await setup_unifi_integration(
+        hass,
+        options={
+            CONF_BLOCK_CLIENT: [BLOCKED["mac"], UNBLOCKED["mac"]],
+            CONF_TRACK_CLIENTS: False,
+            CONF_TRACK_DEVICES: False,
+        },
+        clients_response=[UNBLOCKED],
+        clients_all_response=[BLOCKED],
+    )
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
+
+    blocked = hass.states.get("switch.block_client_1")
+    assert blocked is not None
+    assert blocked.state == "off"
+
+    unblocked = hass.states.get("switch.block_client_2")
+    assert unblocked is not None
+    assert unblocked.state == "on"
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_EVENT},
+        "data": [EVENT_BLOCKED_CLIENT_UNBLOCKED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
+    blocked = hass.states.get("switch.block_client_1")
+    assert blocked is not None
+    assert blocked.state == "on"
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_EVENT},
+        "data": [EVENT_BLOCKED_CLIENT_BLOCKED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
+    blocked = hass.states.get("switch.block_client_1")
+    assert blocked is not None
+    assert blocked.state == "off"
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
+    )
+    assert len(controller.mock_requests) == 5
+    assert controller.mock_requests[4] == {
+        "json": {"mac": "00:00:00:00:01:01", "cmd": "block-sta"},
+        "method": "post",
+        "path": "/cmd/stamgr",
+    }
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
+    )
+    assert len(controller.mock_requests) == 6
+    assert controller.mock_requests[5] == {
+        "json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"},
+        "method": "post",
+        "path": "/cmd/stamgr",
+    }
+
+
 async def test_new_client_discovered_on_block_control(hass):
     """Test if 2nd update has a new client."""
     controller = await setup_unifi_integration(