diff --git a/homeassistant/components/device_tracker/upc_connect.py b/homeassistant/components/device_tracker/upc_connect.py index aafa9824a4e..13336e939a5 100644 --- a/homeassistant/components/device_tracker/upc_connect.py +++ b/homeassistant/components/device_tracker/upc_connect.py @@ -74,8 +74,11 @@ class UPCDeviceScanner(DeviceScanner): return [] raw = yield from self._async_ws_function(CMD_DEVICES) - xml_root = ET.fromstring(raw) + if raw is None: + _LOGGER.warning("Can't read device from %s", self.host) + return + xml_root = ET.fromstring(raw) return [mac.text for mac in xml_root.iter('MACAddr')] @asyncio.coroutine @@ -94,7 +97,8 @@ class UPCDeviceScanner(DeviceScanner): "http://{}/common_page/login.html".format(self.host) ) - self.token = self._async_get_token() + yield from response.text() + self.token = response.cookies['sessionToken'].value # login data = yield from self._async_ws_function(CMD_LOGIN, { @@ -144,7 +148,7 @@ class UPCDeviceScanner(DeviceScanner): # load data, store token for next request raw = yield from response.text() - self.token = self._async_get_token() + self.token = response.cookies['sessionToken'].value return raw @@ -155,10 +159,3 @@ class UPCDeviceScanner(DeviceScanner): finally: if response is not None: yield from response.release() - - def _async_get_token(self): - """Extract token from cookies.""" - cookie_manager = self.websession.cookie_jar.filter_cookies( - "http://{}".format(self.host)) - - return cookie_manager.get('sessionToken') diff --git a/tests/test_util/aiohttp.py b/tests/test_util/aiohttp.py index dcdf69395b4..afe2f626de7 100644 --- a/tests/test_util/aiohttp.py +++ b/tests/test_util/aiohttp.py @@ -37,11 +37,9 @@ class AiohttpClientMocker: content = b'' if params: url = str(yarl.URL(url).with_query(params)) - if cookies: - self._cookies.update(cookies) self._mocks.append(AiohttpClientMockResponse( - method, url, status, content, exc)) + method, url, status, content, cookies, exc)) def get(self, *args, **kwargs): """Register a mock get request.""" @@ -68,10 +66,6 @@ class AiohttpClientMocker: """Number of requests made.""" return len(self.mock_calls) - def filter_cookies(self, host): - """Return hosts cookies.""" - return self._cookies - def clear_requests(self): """Reset mock calls.""" self._mocks.clear() @@ -97,7 +91,7 @@ class AiohttpClientMocker: class AiohttpClientMockResponse: """Mock Aiohttp client response.""" - def __init__(self, method, url, status, response, exc=None): + def __init__(self, method, url, status, response, cookies=None, exc=None): """Initialize a fake response.""" self.method = method self._url = url @@ -107,6 +101,14 @@ class AiohttpClientMockResponse: self.response = response self.exc = exc + self._cookies = {} + + if cookies: + for name, data in cookies.items(): + cookie = mock.MagicMock() + cookie.value = data + self._cookies[name] = cookie + def match_request(self, method, url, params=None): """Test if response answers request.""" if method.lower() != self.method.lower(): @@ -140,6 +142,11 @@ class AiohttpClientMockResponse: return True + @property + def cookies(self): + """Return dict of cookies.""" + return self._cookies + @asyncio.coroutine def read(self): """Return mock response.""" @@ -160,6 +167,10 @@ class AiohttpClientMockResponse: """Mock release.""" pass + def close(self): + """Mock close.""" + pass + @contextmanager def mock_aiohttp_client(): @@ -173,6 +184,4 @@ def mock_aiohttp_client(): setattr(instance, method, functools.partial(mocker.match_request, method)) - instance.cookie_jar.filter_cookies = mocker.filter_cookies - yield mocker