diff --git a/homeassistant/components/onedrive/__init__.py b/homeassistant/components/onedrive/__init__.py index ef7ddd04da6..5feefb2cf7d 100644 --- a/homeassistant/components/onedrive/__init__.py +++ b/homeassistant/components/onedrive/__init__.py @@ -2,8 +2,10 @@ from __future__ import annotations +from collections.abc import Awaitable, Callable from dataclasses import dataclass import logging +from typing import cast from onedrive_personal_sdk import OneDriveClient from onedrive_personal_sdk.exceptions import ( @@ -13,6 +15,7 @@ from onedrive_personal_sdk.exceptions import ( ) from homeassistant.config_entries import ConfigEntry +from homeassistant.const import CONF_ACCESS_TOKEN from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers.aiohttp_client import async_get_clientsession @@ -22,7 +25,6 @@ from homeassistant.helpers.config_entry_oauth2_flow import ( ) from homeassistant.helpers.instance_id import async_get as async_get_instance_id -from .api import OneDriveConfigEntryAccessTokenProvider from .const import DATA_BACKUP_AGENT_LISTENERS, DOMAIN @@ -31,7 +33,7 @@ class OneDriveRuntimeData: """Runtime data for the OneDrive integration.""" client: OneDriveClient - token_provider: OneDriveConfigEntryAccessTokenProvider + token_function: Callable[[], Awaitable[str]] backup_folder_id: str @@ -46,9 +48,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: OneDriveConfigEntry) -> session = OAuth2Session(hass, entry, implementation) - token_provider = OneDriveConfigEntryAccessTokenProvider(session) + async def get_access_token() -> str: + await session.async_ensure_token_valid() + return cast(str, session.token[CONF_ACCESS_TOKEN]) - client = OneDriveClient(token_provider, async_get_clientsession(hass)) + client = OneDriveClient(get_access_token, async_get_clientsession(hass)) # get approot, will be created automatically if it does not exist try: @@ -81,7 +85,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: OneDriveConfigEntry) -> entry.runtime_data = OneDriveRuntimeData( client=client, - token_provider=token_provider, + token_function=get_access_token, backup_folder_id=backup_folder.id, ) diff --git a/homeassistant/components/onedrive/api.py b/homeassistant/components/onedrive/api.py deleted file mode 100644 index d8f6ea188f3..00000000000 --- a/homeassistant/components/onedrive/api.py +++ /dev/null @@ -1,34 +0,0 @@ -"""API for OneDrive bound to Home Assistant OAuth.""" - -from typing import cast - -from onedrive_personal_sdk import TokenProvider - -from homeassistant.const import CONF_ACCESS_TOKEN -from homeassistant.helpers import config_entry_oauth2_flow - - -class OneDriveConfigFlowAccessTokenProvider(TokenProvider): - """Provide OneDrive authentication tied to an OAuth2 based config entry.""" - - def __init__(self, token: str) -> None: - """Initialize OneDrive auth.""" - super().__init__() - self._token = token - - def async_get_access_token(self) -> str: - """Return a valid access token.""" - return self._token - - -class OneDriveConfigEntryAccessTokenProvider(TokenProvider): - """Provide OneDrive authentication tied to an OAuth2 based config entry.""" - - def __init__(self, oauth_session: config_entry_oauth2_flow.OAuth2Session) -> None: - """Initialize OneDrive auth.""" - super().__init__() - self._oauth_session = oauth_session - - def async_get_access_token(self) -> str: - """Return a valid access token.""" - return cast(str, self._oauth_session.token[CONF_ACCESS_TOKEN]) diff --git a/homeassistant/components/onedrive/backup.py b/homeassistant/components/onedrive/backup.py index 43eac020538..78bdcb24b8c 100644 --- a/homeassistant/components/onedrive/backup.py +++ b/homeassistant/components/onedrive/backup.py @@ -109,7 +109,7 @@ class OneDriveBackupAgent(BackupAgent): self._hass = hass self._entry = entry self._client = entry.runtime_data.client - self._token_provider = entry.runtime_data.token_provider + self._token_function = entry.runtime_data.token_function self._folder_id = entry.runtime_data.backup_folder_id self.name = entry.title assert entry.unique_id @@ -145,7 +145,7 @@ class OneDriveBackupAgent(BackupAgent): ) try: item = await LargeFileUploadClient.upload( - self._token_provider, file, session=async_get_clientsession(self._hass) + self._token_function, file, session=async_get_clientsession(self._hass) ) except HashMismatchError as err: raise BackupAgentError( diff --git a/homeassistant/components/onedrive/config_flow.py b/homeassistant/components/onedrive/config_flow.py index cbdf59648b9..900db0177d9 100644 --- a/homeassistant/components/onedrive/config_flow.py +++ b/homeassistant/components/onedrive/config_flow.py @@ -12,7 +12,6 @@ from homeassistant.const import CONF_ACCESS_TOKEN, CONF_TOKEN from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.config_entry_oauth2_flow import AbstractOAuth2FlowHandler -from .api import OneDriveConfigFlowAccessTokenProvider from .const import DOMAIN, OAUTH_SCOPES @@ -36,12 +35,12 @@ class OneDriveConfigFlow(AbstractOAuth2FlowHandler, domain=DOMAIN): data: dict[str, Any], ) -> ConfigFlowResult: """Handle the initial step.""" - token_provider = OneDriveConfigFlowAccessTokenProvider( - cast(str, data[CONF_TOKEN][CONF_ACCESS_TOKEN]) - ) + + async def get_access_token() -> str: + return cast(str, data[CONF_TOKEN][CONF_ACCESS_TOKEN]) graph_client = OneDriveClient( - token_provider, async_get_clientsession(self.hass) + get_access_token, async_get_clientsession(self.hass) ) try: diff --git a/homeassistant/components/onedrive/manifest.json b/homeassistant/components/onedrive/manifest.json index 47eb48742be..88d51e6d73a 100644 --- a/homeassistant/components/onedrive/manifest.json +++ b/homeassistant/components/onedrive/manifest.json @@ -9,5 +9,5 @@ "iot_class": "cloud_polling", "loggers": ["onedrive_personal_sdk"], "quality_scale": "bronze", - "requirements": ["onedrive-personal-sdk==0.0.4"] + "requirements": ["onedrive-personal-sdk==0.0.8"] } diff --git a/requirements_all.txt b/requirements_all.txt index 1254c7521f5..87cb85303ce 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -1559,7 +1559,7 @@ omnilogic==0.4.5 ondilo==0.5.0 # homeassistant.components.onedrive -onedrive-personal-sdk==0.0.4 +onedrive-personal-sdk==0.0.8 # homeassistant.components.onvif onvif-zeep-async==3.2.5 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 758365eeeb7..e7706013267 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -1307,7 +1307,7 @@ omnilogic==0.4.5 ondilo==0.5.0 # homeassistant.components.onedrive -onedrive-personal-sdk==0.0.4 +onedrive-personal-sdk==0.0.8 # homeassistant.components.onvif onvif-zeep-async==3.2.5 diff --git a/tests/components/onedrive/conftest.py b/tests/components/onedrive/conftest.py index e76ce1d01c8..0d6ee09d587 100644 --- a/tests/components/onedrive/conftest.py +++ b/tests/components/onedrive/conftest.py @@ -67,8 +67,8 @@ def mock_config_entry(expires_at: int, scopes: list[str]) -> MockConfigEntry: ) -@pytest.fixture(autouse=True) -def mock_onedrive_client() -> Generator[MagicMock]: +@pytest.fixture +def mock_onedrive_client_init() -> Generator[MagicMock]: """Return a mocked GraphServiceClient.""" with ( patch( @@ -80,19 +80,25 @@ def mock_onedrive_client() -> Generator[MagicMock]: new=onedrive_client, ), ): - client = onedrive_client.return_value - client.get_approot.return_value = MOCK_APPROOT - client.create_folder.return_value = MOCK_BACKUP_FOLDER - client.list_drive_items.return_value = [MOCK_BACKUP_FILE] - client.get_drive_item.return_value = MOCK_BACKUP_FILE + yield onedrive_client - class MockStreamReader: - async def iter_chunked(self, chunk_size: int) -> AsyncIterator[bytes]: - yield b"backup data" - client.download_drive_item.return_value = MockStreamReader() +@pytest.fixture(autouse=True) +def mock_onedrive_client(mock_onedrive_client_init: MagicMock) -> Generator[MagicMock]: + """Return a mocked GraphServiceClient.""" + client = mock_onedrive_client_init.return_value + client.get_approot.return_value = MOCK_APPROOT + client.create_folder.return_value = MOCK_BACKUP_FOLDER + client.list_drive_items.return_value = [MOCK_BACKUP_FILE] + client.get_drive_item.return_value = MOCK_BACKUP_FILE - yield client + class MockStreamReader: + async def iter_chunked(self, chunk_size: int) -> AsyncIterator[bytes]: + yield b"backup data" + + client.download_drive_item.return_value = MockStreamReader() + + return client @pytest.fixture diff --git a/tests/components/onedrive/test_config_flow.py b/tests/components/onedrive/test_config_flow.py index 9acfd8ada3c..fb0d58b86c6 100644 --- a/tests/components/onedrive/test_config_flow.py +++ b/tests/components/onedrive/test_config_flow.py @@ -70,6 +70,7 @@ async def test_full_flow( hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, mock_setup_entry: AsyncMock, + mock_onedrive_client_init: MagicMock, ) -> None: """Check full flow.""" @@ -79,6 +80,10 @@ async def test_full_flow( await _do_get_token(hass, result, hass_client_no_auth, aioclient_mock) result = await hass.config_entries.flow.async_configure(result["flow_id"]) + # Ensure the token callback is set up correctly + token_callback = mock_onedrive_client_init.call_args[0][0] + assert await token_callback() == "mock-access-token" + assert result["type"] is FlowResultType.CREATE_ENTRY assert len(hass.config_entries.async_entries(DOMAIN)) == 1 assert len(mock_setup_entry.mock_calls) == 1 diff --git a/tests/components/onedrive/test_init.py b/tests/components/onedrive/test_init.py index 674bc2d38d9..a6ad55442aa 100644 --- a/tests/components/onedrive/test_init.py +++ b/tests/components/onedrive/test_init.py @@ -16,10 +16,15 @@ from tests.common import MockConfigEntry async def test_load_unload_config_entry( hass: HomeAssistant, mock_config_entry: MockConfigEntry, + mock_onedrive_client_init: MagicMock, ) -> None: """Test loading and unloading the integration.""" await setup_integration(hass, mock_config_entry) + # Ensure the token callback is set up correctly + token_callback = mock_onedrive_client_init.call_args[0][0] + assert await token_callback() == "mock-access-token" + assert mock_config_entry.state is ConfigEntryState.LOADED await hass.config_entries.async_unload(mock_config_entry.entry_id)