diff --git a/homeassistant/components/onedrive/__init__.py b/homeassistant/components/onedrive/__init__.py index 454c782af92..f10b8fe0d91 100644 --- a/homeassistant/components/onedrive/__init__.py +++ b/homeassistant/components/onedrive/__init__.py @@ -41,14 +41,7 @@ _LOGGER = logging.getLogger(__name__) async def async_setup_entry(hass: HomeAssistant, entry: OneDriveConfigEntry) -> bool: """Set up OneDrive from a config entry.""" - implementation = await async_get_config_entry_implementation(hass, entry) - session = OAuth2Session(hass, entry, implementation) - - async def get_access_token() -> str: - await session.async_ensure_token_valid() - return cast(str, session.token[CONF_ACCESS_TOKEN]) - - client = OneDriveClient(get_access_token, async_get_clientsession(hass)) + client, get_access_token = await _get_onedrive_client(hass, entry) # get approot, will be created automatically if it does not exist approot = await _handle_item_operation(client.get_approot, "approot") @@ -164,20 +157,47 @@ async def async_migrate_entry(hass: HomeAssistant, entry: OneDriveConfigEntry) - _LOGGER.debug( "Migrating OneDrive config entry from version %s.%s", version, minor_version ) - + client, _ = await _get_onedrive_client(hass, entry) instance_id = await async_get_instance_id(hass) + try: + approot = await client.get_approot() + folder = await client.get_drive_item( + f"{approot.id}:/backups_{instance_id[:8]}:" + ) + except OneDriveException: + _LOGGER.exception("Migration to version 1.2 failed") + return False + hass.config_entries.async_update_entry( entry, data={ **entry.data, - CONF_FOLDER_ID: "id", # will be updated during setup_entry + CONF_FOLDER_ID: folder.id, CONF_FOLDER_NAME: f"backups_{instance_id[:8]}", }, + minor_version=2, ) _LOGGER.debug("Migration to version 1.2 successful") return True +async def _get_onedrive_client( + hass: HomeAssistant, entry: OneDriveConfigEntry +) -> tuple[OneDriveClient, Callable[[], Awaitable[str]]]: + """Get OneDrive client.""" + implementation = await async_get_config_entry_implementation(hass, entry) + session = OAuth2Session(hass, entry, implementation) + + async def get_access_token() -> str: + await session.async_ensure_token_valid() + return cast(str, session.token[CONF_ACCESS_TOKEN]) + + return ( + OneDriveClient(get_access_token, async_get_clientsession(hass)), + get_access_token, + ) + + async def _handle_item_operation( func: Callable[[], Awaitable[Item]], folder: str ) -> Item: diff --git a/tests/components/onedrive/test_init.py b/tests/components/onedrive/test_init.py index c7765e0a7f8..952ca01e1cb 100644 --- a/tests/components/onedrive/test_init.py +++ b/tests/components/onedrive/test_init.py @@ -236,7 +236,6 @@ async def test_data_cap_issues( async def test_1_1_to_1_2_migration( hass: HomeAssistant, - mock_onedrive_client: MagicMock, mock_config_entry: MockConfigEntry, mock_folder: Folder, ) -> None: @@ -251,12 +250,34 @@ async def test_1_1_to_1_2_migration( }, ) + await setup_integration(hass, old_config_entry) + assert old_config_entry.data[CONF_FOLDER_ID] == mock_folder.id + assert old_config_entry.data[CONF_FOLDER_NAME] == mock_folder.name + assert old_config_entry.minor_version == 2 + + +async def test_1_1_to_1_2_migration_failure( + hass: HomeAssistant, + mock_onedrive_client: MagicMock, + mock_config_entry: MockConfigEntry, +) -> None: + """Test migration from 1.1 to 1.2 failure.""" + old_config_entry = MockConfigEntry( + unique_id="mock_drive_id", + title="John Doe's OneDrive", + domain=DOMAIN, + data={ + "auth_implementation": mock_config_entry.data["auth_implementation"], + "token": mock_config_entry.data["token"], + }, + ) + # will always 404 after migration, because of dummy id mock_onedrive_client.get_drive_item.side_effect = NotFoundError(404, "Not found") await setup_integration(hass, old_config_entry) - assert old_config_entry.data[CONF_FOLDER_ID] == mock_folder.id - assert old_config_entry.data[CONF_FOLDER_NAME] == mock_folder.name + assert old_config_entry.state is ConfigEntryState.MIGRATION_ERROR + assert old_config_entry.minor_version == 1 async def test_migration_guard_against_major_downgrade(