From 2b4980ae5d43b75fffe2e4955639679d74aab66d Mon Sep 17 00:00:00 2001
From: Dan Ford <dpford@gmail.com>
Date: Sun, 21 Aug 2016 13:09:44 -0700
Subject: [PATCH] Add tplink Archer C7 device tracking support for 5Ghz
 networks

---
 .../components/device_tracker/tplink.py       | 26 ++++---
 .../components/device_tracker/test_tplink.py  | 69 +++++++++++++++++++
 2 files changed, 84 insertions(+), 11 deletions(-)
 create mode 100644 tests/components/device_tracker/test_tplink.py

diff --git a/homeassistant/components/device_tracker/tplink.py b/homeassistant/components/device_tracker/tplink.py
index 37d4d2b4471..17beab02532 100755
--- a/homeassistant/components/device_tracker/tplink.py
+++ b/homeassistant/components/device_tracker/tplink.py
@@ -313,19 +313,23 @@ class Tplink4DeviceScanner(TplinkDeviceScanner):
 
             _LOGGER.info("Loading wireless clients...")
 
-            url = 'http://{}/{}/userRpm/WlanStationRpm.htm' \
-                .format(self.host, self.token)
-            referer = 'http://{}'.format(self.host)
-            cookie = 'Authorization=Basic {}'.format(self.credentials)
+            mac_results = []
 
-            page = requests.get(url, headers={
-                'cookie': cookie,
-                'referer': referer
-            })
-            result = self.parse_macs.findall(page.text)
+            # Check both the 2.4GHz and 5GHz client list URLs
+            for clients_url in ('WlanStationRpm.htm', 'WlanStationRpm_5g.htm'):
+                url = 'http://{}/{}/userRpm/{}' \
+                    .format(self.host, self.token, clients_url)
+                referer = 'http://{}'.format(self.host)
+                cookie = 'Authorization=Basic {}'.format(self.credentials)
 
-            if not result:
+                page = requests.get(url, headers={
+                    'cookie': cookie,
+                    'referer': referer
+                })
+                mac_results.extend(self.parse_macs.findall(page.text))
+
+            if not mac_results:
                 return False
 
-            self.last_results = [mac.replace("-", ":") for mac in result]
+            self.last_results = [mac.replace("-", ":") for mac in mac_results]
             return True
diff --git a/tests/components/device_tracker/test_tplink.py b/tests/components/device_tracker/test_tplink.py
new file mode 100644
index 00000000000..715d4f3ffde
--- /dev/null
+++ b/tests/components/device_tracker/test_tplink.py
@@ -0,0 +1,69 @@
+"""The tests for the tplink device tracker platform."""
+
+import os
+import unittest
+
+from homeassistant.components import device_tracker
+from homeassistant.components.device_tracker.tplink import Tplink4DeviceScanner
+from homeassistant.const import (CONF_PLATFORM, CONF_PASSWORD, CONF_USERNAME,
+                                 CONF_HOST)
+import requests_mock
+
+from tests.common import get_test_home_assistant
+
+
+class TestTplink4DeviceScanner(unittest.TestCase):
+    """Tests for the Tplink4DeviceScanner class."""
+
+    def setUp(self):  # pylint: disable=invalid-name
+        """Setup things to be run when tests are started."""
+        self.hass = get_test_home_assistant()
+
+    def tearDown(self):  # pylint: disable=invalid-name
+        """Stop everything that was started."""
+        try:
+            os.remove(self.hass.config.path(device_tracker.YAML_DEVICES))
+        except FileNotFoundError:
+            pass
+
+    @requests_mock.mock()
+    def test_get_mac_addresses_from_both_bands(self, m):
+        """
+        Test grabbing the mac addresses from both 2.4 and 5 GHz clients pages.
+        """
+        conf_dict = {
+            CONF_PLATFORM: 'tplink',
+            CONF_HOST: 'fake_host',
+            CONF_USERNAME: 'fake_user',
+            CONF_PASSWORD: 'fake_pass'
+        }
+
+        # Mock the token retrieval process
+        FAKE_TOKEN = 'fake_token'
+        fake_auth_token_response = 'window.parent.location.href = ' \
+                                   '"https://a/{}/userRpm/Index.htm";'.format(
+                                       FAKE_TOKEN)
+
+        m.get('http://{}/userRpm/LoginRpm.htm?Save=Save'.format(
+            conf_dict[CONF_HOST]), text=fake_auth_token_response)
+
+        FAKE_MAC_1 = 'CA-FC-8A-C8-BB-53'
+        FAKE_MAC_2 = '6C-48-83-21-46-8D'
+        FAKE_MAC_3 = '77-98-75-65-B1-2B'
+        mac_response_2_4 = '{} {}'.format(FAKE_MAC_1, FAKE_MAC_2)
+        mac_response_5 = '{}'.format(FAKE_MAC_3)
+
+        # Mock the 2.4 GHz clients page
+        m.get('http://{}/{}/userRpm/WlanStationRpm.htm'.format(
+            conf_dict[CONF_HOST], FAKE_TOKEN), text=mac_response_2_4)
+
+        # Mock the 5 GHz clients page
+        m.get('http://{}/{}/userRpm/WlanStationRpm_5g.htm'.format(
+            conf_dict[CONF_HOST], FAKE_TOKEN), text=mac_response_5)
+
+        tplink = Tplink4DeviceScanner(conf_dict)
+
+        expected_mac_results = [mac.replace('-', ':') for mac in
+                                [FAKE_MAC_1, FAKE_MAC_2, FAKE_MAC_3]]
+
+        self.assertEquals(tplink.last_results, expected_mac_results)