From f72d9a2c023a1e1c41065cc56914ab3caa3165a2 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 15 Aug 2024 14:46:23 +0200
Subject: [PATCH] Raise on database error in recorder.migration._modify_columns
 (#123642)

* Raise on database error in recorder.migration._modify_columns

* Improve test coverage
---
 homeassistant/components/recorder/migration.py |  1 +
 tests/components/recorder/conftest.py          |  3 +++
 tests/components/recorder/test_migrate.py      | 11 ++++++++---
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/recorder/migration.py b/homeassistant/components/recorder/migration.py
index 44b4a980238..9a19ff7084a 100644
--- a/homeassistant/components/recorder/migration.py
+++ b/homeassistant/components/recorder/migration.py
@@ -580,6 +580,7 @@ def _modify_columns(
                 _LOGGER.exception(
                     "Could not modify column %s in table %s", column_def, table_name
                 )
+                raise
 
 
 def _update_states_table_with_foreign_key_options(
diff --git a/tests/components/recorder/conftest.py b/tests/components/recorder/conftest.py
index b13b3b270a9..9cdf9dbb372 100644
--- a/tests/components/recorder/conftest.py
+++ b/tests/components/recorder/conftest.py
@@ -70,6 +70,7 @@ class InstrumentedMigration:
     apply_update_mock: Mock
     stall_on_schema_version: int | None
     apply_update_stalled: threading.Event
+    apply_update_version: int | None
 
 
 @pytest.fixture(name="instrument_migration")
@@ -147,6 +148,7 @@ def instrument_migration(
         old_version: int,
     ):
         """Control migration progress."""
+        instrumented_migration.apply_update_version = new_version
         stall_version = instrumented_migration.stall_on_schema_version
         if stall_version is None or stall_version == new_version:
             instrumented_migration.apply_update_stalled.set()
@@ -182,6 +184,7 @@ def instrument_migration(
             apply_update_mock=apply_update_mock,
             stall_on_schema_version=None,
             apply_update_stalled=threading.Event(),
+            apply_update_version=None,
         )
 
         instrumented_migration.live_migration_done_stall.set()
diff --git a/tests/components/recorder/test_migrate.py b/tests/components/recorder/test_migrate.py
index 7b91a97e5a0..4bc317bdaa7 100644
--- a/tests/components/recorder/test_migrate.py
+++ b/tests/components/recorder/test_migrate.py
@@ -208,8 +208,12 @@ async def test_database_migration_failed(
         "expected_pn_dismiss",
     ),
     [
-        (11, "DropConstraint", False, 1, 0),
-        (44, "DropConstraint", False, 2, 1),
+        # Test error handling in _update_states_table_with_foreign_key_options
+        (11, "homeassistant.components.recorder.migration.DropConstraint", False, 1, 0),
+        # Test error handling in _modify_columns
+        (12, "sqlalchemy.engine.base.Connection.execute", False, 1, 0),
+        # Test error handling in _drop_foreign_key_constraints
+        (44, "homeassistant.components.recorder.migration.DropConstraint", False, 2, 1),
     ],
 )
 @pytest.mark.skip_on_db_engine(["sqlite"])
@@ -255,7 +259,7 @@ async def test_database_migration_failed_non_sqlite(
 
         # Make it fail
         with patch(
-            f"homeassistant.components.recorder.migration.{func_to_patch}",
+            func_to_patch,
             side_effect=OperationalError(
                 None, None, OSError("No space left on device")
             ),
@@ -267,6 +271,7 @@ async def test_database_migration_failed_non_sqlite(
             await hass.async_add_executor_job(recorder.get_instance(hass).join)
             await hass.async_block_till_done()
 
+    assert instrument_migration.apply_update_version == patch_version
     assert recorder.util.async_migration_in_progress(hass) is False
     assert len(mock_create.mock_calls) == expected_pn_create
     assert len(mock_dismiss.mock_calls) == expected_pn_dismiss