From c3bb5417ffa81890a89e53dd9ef4919ffb45778e Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sun, 19 Jun 2022 22:30:03 -0700 Subject: [PATCH 1/4] Rewrite google calendar oauth exchange callbacks --- homeassistant/components/google/api.py | 96 +++++++++++-------- .../components/google/config_flow.py | 8 +- tests/components/google/test_config_flow.py | 15 +-- 3 files changed, 71 insertions(+), 48 deletions(-) diff --git a/homeassistant/components/google/api.py b/homeassistant/components/google/api.py index f4a4912a1b9675..2e9de44e28124f 100644 --- a/homeassistant/components/google/api.py +++ b/homeassistant/components/google/api.py @@ -2,7 +2,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable import datetime import logging from typing import Any, cast @@ -19,9 +18,12 @@ from homeassistant.components.application_credentials import AuthImplementation from homeassistant.config_entries import ConfigEntry -from homeassistant.core import CALLBACK_TYPE, HomeAssistant +from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback from homeassistant.helpers import config_entry_oauth2_flow -from homeassistant.helpers.event import async_track_time_interval +from homeassistant.helpers.event import ( + async_track_point_in_utc_time, + async_track_time_interval, +) from homeassistant.util import dt from .const import ( @@ -76,6 +78,16 @@ def __init__( self._oauth_flow = oauth_flow self._device_flow_info: DeviceFlowInfo = device_flow_info self._exchange_task_unsub: CALLBACK_TYPE | None = None + self._timeout_unsub: CALLBACK_TYPE | None = None + max_timeout = dt.utcnow() + datetime.timedelta(seconds=EXCHANGE_TIMEOUT_SECONDS) + # For some reason, oauth.step1_get_device_and_user_codes() returns a datetime + # object without tzinfo. For the comparison below to work, it needs one. + user_code_expiry = device_flow_info.user_code_expiry.replace( + tzinfo=datetime.timezone.utc + ) + self._expiration_time = min(user_code_expiry, max_timeout) + self._listener: CALLBACK_TYPE | None = None + self._creds: Credentials | None = None @property def verification_url(self) -> str: @@ -87,48 +99,56 @@ def user_code(self) -> str: """Return the code that the user should enter at the verification url.""" return self._device_flow_info.user_code # type: ignore[no-any-return] - async def start_exchange_task( - self, finished_cb: Callable[[Credentials | None], Awaitable[None]] + @callback + def async_set_listener( + self, + update_callback: CALLBACK_TYPE, ) -> None: - """Start the device auth exchange flow polling. + """Invoke the update callback when the exchange finishes or on timeout.""" + self._listener = update_callback - The callback is invoked with the valid credentials or with None on timeout. - """ - _LOGGER.debug("Starting exchange flow") - assert not self._exchange_task_unsub - max_timeout = dt.utcnow() + datetime.timedelta(seconds=EXCHANGE_TIMEOUT_SECONDS) - # For some reason, oauth.step1_get_device_and_user_codes() returns a datetime - # object without tzinfo. For the comparison below to work, it needs one. - user_code_expiry = self._device_flow_info.user_code_expiry.replace( - tzinfo=datetime.timezone.utc - ) - expiration_time = min(user_code_expiry, max_timeout) - - def _exchange() -> Credentials: - return self._oauth_flow.step2_exchange( - device_flow_info=self._device_flow_info - ) - - async def _poll_attempt(now: datetime.datetime) -> None: - assert self._exchange_task_unsub - _LOGGER.debug("Attempting OAuth code exchange") - # Note: The callback is invoked with None when the device code has expired - creds: Credentials | None = None - if now < expiration_time: - try: - creds = await self._hass.async_add_executor_job(_exchange) - except FlowExchangeError: - _LOGGER.debug("Token not yet ready; trying again later") - return - self._exchange_task_unsub() - self._exchange_task_unsub = None - await finished_cb(creds) + @property + def creds(self) -> Credentials | None: + """Return result of exchange step or None on timeout.""" + return self._creds + def async_start_exchange(self) -> None: + """Start the device auth exchange flow polling.""" + _LOGGER.debug("Starting exchange flow") self._exchange_task_unsub = async_track_time_interval( self._hass, - _poll_attempt, + self._async_poll_attempt, datetime.timedelta(seconds=self._device_flow_info.interval), ) + self._timeout_unsub = async_track_point_in_utc_time( + self._hass, self._async_timeout, self._expiration_time + ) + + async def _async_poll_attempt(self, now: datetime.datetime) -> None: + _LOGGER.debug("Attempting OAuth code exchange") + try: + self._creds = await self._hass.async_add_executor_job(self._exchange) + except FlowExchangeError: + _LOGGER.debug("Token not yet ready; trying again later") + return + self._finish() + + def _exchange(self) -> Credentials: + return self._oauth_flow.step2_exchange(device_flow_info=self._device_flow_info) + + @callback + def _async_timeout(self, now: datetime.datetime) -> None: + _LOGGER.debug("OAuth token exchange timeout.") + self._finish() + + @callback + def _finish(self) -> None: + if self._exchange_task_unsub: + self._exchange_task_unsub() + if self._timeout_unsub: + self._timeout_unsub() + if self._listener: + self._listener() def get_feature_access( diff --git a/homeassistant/components/google/config_flow.py b/homeassistant/components/google/config_flow.py index 046840075ff3ba..609b85963d953b 100644 --- a/homeassistant/components/google/config_flow.py +++ b/homeassistant/components/google/config_flow.py @@ -6,7 +6,6 @@ from gcal_sync.api import GoogleCalendarService from gcal_sync.exceptions import ApiException -from oauth2client.client import Credentials import voluptuous as vol from homeassistant import config_entries @@ -96,9 +95,9 @@ async def async_step_auth( return self.async_abort(reason="oauth_error") self._device_flow = device_flow - async def _exchange_finished(creds: Credentials | None) -> None: + def _exchange_finished() -> None: self.external_data = { - DEVICE_AUTH_CREDS: creds + DEVICE_AUTH_CREDS: device_flow.creds } # is None on timeout/expiration self.hass.async_create_task( self.hass.config_entries.flow.async_configure( @@ -106,7 +105,8 @@ async def _exchange_finished(creds: Credentials | None) -> None: ) ) - await device_flow.start_exchange_task(_exchange_finished) + device_flow.async_set_listener(_exchange_finished) + device_flow.async_start_exchange() return self.async_show_progress( step_id="auth", diff --git a/tests/components/google/test_config_flow.py b/tests/components/google/test_config_flow.py index 00f50e129e4f90..78a61da90edb92 100644 --- a/tests/components/google/test_config_flow.py +++ b/tests/components/google/test_config_flow.py @@ -248,13 +248,13 @@ async def test_code_error( assert result.get("reason") == "oauth_error" -@pytest.mark.parametrize("code_expiration_delta", [datetime.timedelta(minutes=-5)]) +@pytest.mark.parametrize("code_expiration_delta", [datetime.timedelta(seconds=5)]) async def test_expired_after_exchange( hass: HomeAssistant, mock_code_flow: Mock, component_setup: ComponentSetup, ) -> None: - """Test successful creds setup.""" + """Test credential exchange expires immediately.""" assert await component_setup() result = await hass.config_entries.flow.async_init( @@ -265,10 +265,13 @@ async def test_expired_after_exchange( assert "description_placeholders" in result assert "url" in result["description_placeholders"] - # Run one tick to invoke the credential exchange check - now = utcnow() - await fire_alarm(hass, now + CODE_CHECK_ALARM_TIMEDELTA) - await hass.async_block_till_done() + with patch( + "oauth2client.client.OAuth2WebServerFlow.step2_exchange", + side_effect=FlowExchangeError(), + ): + now = utcnow() + await fire_alarm(hass, now + datetime.timedelta(seconds=65)) + await hass.async_block_till_done() result = await hass.config_entries.flow.async_configure(flow_id=result["flow_id"]) assert result.get("type") == "abort" From a8e3c1a4ba1670b85f114431a4e008ff3ea71066 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sun, 19 Jun 2022 22:35:22 -0700 Subject: [PATCH 2/4] Update config flow test wording --- tests/components/google/test_config_flow.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/components/google/test_config_flow.py b/tests/components/google/test_config_flow.py index 78a61da90edb92..3d18faae47169c 100644 --- a/tests/components/google/test_config_flow.py +++ b/tests/components/google/test_config_flow.py @@ -248,13 +248,13 @@ async def test_code_error( assert result.get("reason") == "oauth_error" -@pytest.mark.parametrize("code_expiration_delta", [datetime.timedelta(seconds=5)]) +@pytest.mark.parametrize("code_expiration_delta", [datetime.timedelta(seconds=50)]) async def test_expired_after_exchange( hass: HomeAssistant, mock_code_flow: Mock, component_setup: ComponentSetup, ) -> None: - """Test credential exchange expires immediately.""" + """Test credential exchange expires.""" assert await component_setup() result = await hass.config_entries.flow.async_init( @@ -265,6 +265,7 @@ async def test_expired_after_exchange( assert "description_placeholders" in result assert "url" in result["description_placeholders"] + # Fail first attempt then advance clock past exchange timeout with patch( "oauth2client.client.OAuth2WebServerFlow.step2_exchange", side_effect=FlowExchangeError(), From 0b24661280417e826ee820713073696990b36297 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Fri, 24 Jun 2022 07:13:29 -0700 Subject: [PATCH 3/4] Remove unnecessary period in log message --- homeassistant/components/google/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/homeassistant/components/google/api.py b/homeassistant/components/google/api.py index 2e9de44e28124f..4ed39d1cc146e5 100644 --- a/homeassistant/components/google/api.py +++ b/homeassistant/components/google/api.py @@ -138,7 +138,7 @@ def _exchange(self) -> Credentials: @callback def _async_timeout(self, now: datetime.datetime) -> None: - _LOGGER.debug("OAuth token exchange timeout.") + _LOGGER.debug("OAuth token exchange timeout") self._finish() @callback From 063da908476c323a3f44f753b34bca2177ac6a96 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Fri, 24 Jun 2022 08:27:27 -0700 Subject: [PATCH 4/4] Fix issues with tests using freezetime that doesn't work with timezones/utc --- homeassistant/components/google/api.py | 17 +++++----- tests/components/google/test_config_flow.py | 36 +++++++++++++-------- 2 files changed, 31 insertions(+), 22 deletions(-) diff --git a/homeassistant/components/google/api.py b/homeassistant/components/google/api.py index 4ed39d1cc146e5..47aa32dcd11cde 100644 --- a/homeassistant/components/google/api.py +++ b/homeassistant/components/google/api.py @@ -79,13 +79,6 @@ def __init__( self._device_flow_info: DeviceFlowInfo = device_flow_info self._exchange_task_unsub: CALLBACK_TYPE | None = None self._timeout_unsub: CALLBACK_TYPE | None = None - max_timeout = dt.utcnow() + datetime.timedelta(seconds=EXCHANGE_TIMEOUT_SECONDS) - # For some reason, oauth.step1_get_device_and_user_codes() returns a datetime - # object without tzinfo. For the comparison below to work, it needs one. - user_code_expiry = device_flow_info.user_code_expiry.replace( - tzinfo=datetime.timezone.utc - ) - self._expiration_time = min(user_code_expiry, max_timeout) self._listener: CALLBACK_TYPE | None = None self._creds: Credentials | None = None @@ -115,13 +108,21 @@ def creds(self) -> Credentials | None: def async_start_exchange(self) -> None: """Start the device auth exchange flow polling.""" _LOGGER.debug("Starting exchange flow") + max_timeout = dt.utcnow() + datetime.timedelta(seconds=EXCHANGE_TIMEOUT_SECONDS) + # For some reason, oauth.step1_get_device_and_user_codes() returns a datetime + # object without tzinfo. For the comparison below to work, it needs one. + user_code_expiry = self._device_flow_info.user_code_expiry.replace( + tzinfo=datetime.timezone.utc + ) + expiration_time = min(user_code_expiry, max_timeout) + self._exchange_task_unsub = async_track_time_interval( self._hass, self._async_poll_attempt, datetime.timedelta(seconds=self._device_flow_info.interval), ) self._timeout_unsub = async_track_point_in_utc_time( - self._hass, self._async_timeout, self._expiration_time + self._hass, self._async_timeout, expiration_time ) async def _async_poll_attempt(self, now: datetime.datetime) -> None: diff --git a/tests/components/google/test_config_flow.py b/tests/components/google/test_config_flow.py index 3d18faae47169c..24ad8a7b76911a 100644 --- a/tests/components/google/test_config_flow.py +++ b/tests/components/google/test_config_flow.py @@ -10,6 +10,7 @@ from aiohttp.client_exceptions import ClientError from freezegun.api import FrozenDateTimeFactory from oauth2client.client import ( + DeviceFlowInfo, FlowExchangeError, OAuth2Credentials, OAuth2DeviceCodeError, @@ -59,10 +60,17 @@ async def mock_code_flow( ) -> YieldFixture[Mock]: """Fixture for initiating OAuth flow.""" with patch( - "oauth2client.client.OAuth2WebServerFlow.step1_get_device_and_user_codes", + "homeassistant.components.google.api.OAuth2WebServerFlow.step1_get_device_and_user_codes", ) as mock_flow: - mock_flow.return_value.user_code_expiry = utcnow() + code_expiration_delta - mock_flow.return_value.interval = CODE_CHECK_INTERVAL + mock_flow.return_value = DeviceFlowInfo.FromResponse( + { + "device_code": "4/4-GMMhmHCXhWEzkobqIHGG_EnNYYsAkukHspeYUk9E8", + "user_code": "GQVQ-JKEC", + "verification_url": "https://www.google.com/device", + "expires_in": code_expiration_delta.total_seconds(), + "interval": CODE_CHECK_INTERVAL, + } + ) yield mock_flow @@ -70,7 +78,8 @@ async def mock_code_flow( async def mock_exchange(creds: OAuth2Credentials) -> YieldFixture[Mock]: """Fixture for mocking out the exchange for credentials.""" with patch( - "oauth2client.client.OAuth2WebServerFlow.step2_exchange", return_value=creds + "homeassistant.components.google.api.OAuth2WebServerFlow.step2_exchange", + return_value=creds, ) as mock: yield mock @@ -108,7 +117,6 @@ async def fire_alarm(hass, point_in_time): await hass.async_block_till_done() -@pytest.mark.freeze_time("2022-06-03 15:19:59-00:00") async def test_full_flow_yaml_creds( hass: HomeAssistant, mock_code_flow: Mock, @@ -131,9 +139,8 @@ async def test_full_flow_yaml_creds( "homeassistant.components.google.async_setup_entry", return_value=True ) as mock_setup: # Run one tick to invoke the credential exchange check - freezer.tick(CODE_CHECK_ALARM_TIMEDELTA) - await fire_alarm(hass, datetime.datetime.utcnow()) - await hass.async_block_till_done() + now = utcnow() + await fire_alarm(hass, now + CODE_CHECK_ALARM_TIMEDELTA) result = await hass.config_entries.flow.async_configure( flow_id=result["flow_id"] ) @@ -143,11 +150,12 @@ async def test_full_flow_yaml_creds( assert "data" in result data = result["data"] assert "token" in data + assert 0 < data["token"]["expires_in"] <= 60 * 60 assert ( - data["token"]["expires_in"] - == 60 * 60 - CODE_CHECK_ALARM_TIMEDELTA.total_seconds() + datetime.datetime.now().timestamp() + <= data["token"]["expires_at"] + < (datetime.datetime.now() + datetime.timedelta(days=8)).timestamp() ) - assert data["token"]["expires_at"] == 1654273199.0 data["token"].pop("expires_at") data["token"].pop("expires_in") assert data == { @@ -238,7 +246,7 @@ async def test_code_error( assert await component_setup() with patch( - "oauth2client.client.OAuth2WebServerFlow.step1_get_device_and_user_codes", + "homeassistant.components.google.api.OAuth2WebServerFlow.step1_get_device_and_user_codes", side_effect=OAuth2DeviceCodeError("Test Failure"), ): result = await hass.config_entries.flow.async_init( @@ -267,7 +275,7 @@ async def test_expired_after_exchange( # Fail first attempt then advance clock past exchange timeout with patch( - "oauth2client.client.OAuth2WebServerFlow.step2_exchange", + "homeassistant.components.google.api.OAuth2WebServerFlow.step2_exchange", side_effect=FlowExchangeError(), ): now = utcnow() @@ -299,7 +307,7 @@ async def test_exchange_error( # Run one tick to invoke the credential exchange check now = utcnow() with patch( - "oauth2client.client.OAuth2WebServerFlow.step2_exchange", + "homeassistant.components.google.api.OAuth2WebServerFlow.step2_exchange", side_effect=FlowExchangeError(), ): now += CODE_CHECK_ALARM_TIMEDELTA