Improve onedrive migration (#139458)

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)
+    client, get_access_token = await _get_onedrive_client(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))
 
     # 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:

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.state is ConfigEntryState.MIGRATION_ERROR
-    assert old_config_entry.data[CONF_FOLDER_NAME] == mock_folder.name
+    assert old_config_entry.minor_version == 1
 
 
 async def test_migration_guard_against_major_downgrade(