zigpy 1.1.1__tar.gz → 1.2.1__tar.gz

{zigpy-1.1.1/zigpy.egg-info → zigpy-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zigpy
-Version: 1.1.1
+Version: 1.2.1
 Summary: Library implementing a Zigbee stack
 Author-email: Russell Cloran <rcloran@gmail.com>
 License: GPL-3.0

{zigpy-1.1.1 → zigpy-1.2.1}/tests/test_appdb.py

@@ -1,8 +1,6 @@
 import asyncio
-import contextlib
 from datetime import UTC, datetime, timedelta
 import pathlib
-import threading
 import time
 
 import aiosqlite
@@ -42,30 +40,17 @@ from zigpy.quirks.registry import DeviceRegistry
 from zigpy.quirks.v2 import QuirkBuilder
 import zigpy.types as t
 import zigpy.zcl
-from zigpy.zcl import ClusterType, UnsupportedAttribute
+from zigpy.zcl import (
+    ClusterType,
+    OtaQueryCacheClearedEvent,
+    OtaQueryCacheUpdatedEvent,
+    UnsupportedAttribute,
+)
 from zigpy.zcl.clusters.general import Basic, Identify, OnOff, Ota
 from zigpy.zcl.foundation import Status as ZCLStatus, ZCLAttributeDef
 from zigpy.zdo import types as zdo_t
 
-
-@pytest.fixture(autouse=True)
-def auto_kill_aiosqlite():
-    """Aiosqlite's background thread does not let pytest exit when a failure occurs"""
-    yield
-
-    for thread in threading.enumerate():
-        if not isinstance(thread, aiosqlite.core.Connection):
-            continue
-
-        try:
-            conn = thread._conn
-        except ValueError:
-            pass
-        else:
-            with contextlib.suppress(zigpy.appdb.sqlite3.ProgrammingError):
-                conn.close()
-
-        thread._running = False
+pytestmark = pytest.mark.usefixtures("auto_kill_aiosqlite")
 
 
 async def make_app_with_db(database_file):
@@ -1607,3 +1592,231 @@ async def test_device_signature_ignores_quirks(tmp_path) -> None:
     assert dev2.original_signature == expected_signature
 
     await app2.shutdown()
+
+
+@patch("zigpy.device.Device.schedule_initialize", new=mock_dev_init(True))
+async def test_ota_query_cache_persistence(tmp_path):
+    """Test that OTA query cache is persisted and restored from the database."""
+    db = tmp_path / "test.db"
+    app = await make_app_with_db(db)
+    ieee = make_ieee()
+    app.handle_join(99, ieee, 0)
+
+    dev = app.get_device(ieee)
+    ep = dev.add_endpoint(1)
+    ep.status = zigpy.endpoint.Status.ZDO_INIT
+    ep.profile_id = 260
+    ep.device_type = profiles.zha.DeviceType.PUMP
+    ep.add_input_cluster(0)  # Basic cluster, exercises non-OTA skip in load
+    ota_cluster = ep.add_output_cluster(Ota.cluster_id)
+    app.device_initialized(dev)
+
+    # With hardware_version
+    cmd = Ota.QueryNextImageCommand(
+        field_control=Ota.QueryNextImageCommand.FieldControl.HardwareVersion,
+        manufacturer_code=0x1234,
+        image_type=0x5678,
+        current_file_version=0x000A0001,
+    )
+    cmd.hardware_version = 3
+    ota_cluster.last_query_cmd = cmd
+
+    app.device_initialized(dev)
+    await app.shutdown()
+
+    app2 = await make_app_with_db(db)
+    dev2 = app2.get_device(ieee)
+    ota2 = dev2.endpoints[1].out_clusters[Ota.cluster_id]
+
+    assert ota2.last_query_cmd is not None
+    assert ota2.last_query_cmd.manufacturer_code == 0x1234
+    assert ota2.last_query_cmd.image_type == 0x5678
+    assert ota2.last_query_cmd.current_file_version == 0x000A0001
+    assert ota2.last_query_cmd.hardware_version == 3
+    assert dev2.get_last_ota_query_cmd() is ota2.last_query_cmd
+
+    # Update to a command without hardware_version
+    ota2.last_query_cmd = Ota.QueryNextImageCommand(
+        field_control=Ota.QueryNextImageCommand.FieldControl(0),
+        manufacturer_code=0xAAAA,
+        image_type=0xBBBB,
+        current_file_version=0x00000042,
+    )
+    app2.device_initialized(dev2)
+    await app2.shutdown()
+
+    app3 = await make_app_with_db(db)
+    dev3 = app3.get_device(ieee)
+    ota3 = dev3.endpoints[1].out_clusters[Ota.cluster_id]
+
+    assert ota3.last_query_cmd is not None
+    assert ota3.last_query_cmd.manufacturer_code == 0xAAAA
+    assert ota3.last_query_cmd.current_file_version == 0x00000042
+    assert not hasattr(ota3.last_query_cmd, "hardware_version") or (
+        ota3.last_query_cmd.hardware_version is None
+    )
+
+    await app3.shutdown()
+
+
+@patch("zigpy.device.Device.schedule_initialize", new=mock_dev_init(True))
+async def test_ota_query_cache_event_save(tmp_path):
+    """Test that the OtaQueryCacheUpdatedEvent triggers a DB save."""
+    db = tmp_path / "test.db"
+    app = await make_app_with_db(db)
+    ieee = make_ieee()
+    app.handle_join(99, ieee, 0)
+
+    dev = app.get_device(ieee)
+    ep = dev.add_endpoint(1)
+    ep.status = zigpy.endpoint.Status.ZDO_INIT
+    ep.profile_id = 260
+    ep.device_type = profiles.zha.DeviceType.PUMP
+    ota_cluster = ep.add_output_cluster(Ota.cluster_id)
+    app.device_initialized(dev)
+
+    # Simulate the event that _handle_query_next_image would emit
+    event = OtaQueryCacheUpdatedEvent(
+        device_ieee=str(ieee),
+        endpoint_id=1,
+        cluster_type=ota_cluster.cluster_type,
+        cluster_id=ota_cluster.cluster_id,
+        manufacturer_code=0x1111,
+        image_type=0x2222,
+        current_file_version=0x00000099,
+        hardware_version=7,
+    )
+
+    ota_cluster.emit(OtaQueryCacheUpdatedEvent.event_type, event)
+    await app.shutdown()
+
+    app2 = await make_app_with_db(db)
+    dev2 = app2.get_device(ieee)
+    ota2 = dev2.endpoints[1].out_clusters[Ota.cluster_id]
+
+    assert ota2.last_query_cmd is not None
+    assert ota2.last_query_cmd.manufacturer_code == 0x1111
+    assert ota2.last_query_cmd.image_type == 0x2222
+    assert ota2.last_query_cmd.current_file_version == 0x00000099
+    assert ota2.last_query_cmd.hardware_version == 7
+
+    await app2.shutdown()
+
+
+@patch("zigpy.device.Device.schedule_initialize", new=mock_dev_init(True))
+async def test_ota_query_cache_cleared_after_update(tmp_path):
+    """Test that OTA query cache is deleted from DB when cleared after an update."""
+    db = tmp_path / "test.db"
+    app = await make_app_with_db(db)
+    ieee = make_ieee()
+    app.handle_join(99, ieee, 0)
+
+    dev = app.get_device(ieee)
+    ep = dev.add_endpoint(1)
+    ep.status = zigpy.endpoint.Status.ZDO_INIT
+    ep.profile_id = 260
+    ep.device_type = profiles.zha.DeviceType.PUMP
+    ota_cluster = ep.add_output_cluster(Ota.cluster_id)
+    app.device_initialized(dev)
+
+    # Save a query cmd
+    ota_cluster.last_query_cmd = Ota.QueryNextImageCommand(
+        field_control=Ota.QueryNextImageCommand.FieldControl(0),
+        manufacturer_code=0x1234,
+        image_type=0x5678,
+        current_file_version=0x00000001,
+    )
+    app.device_initialized(dev)
+
+    # Clear the cache (as update_firmware does after a successful OTA)
+    ota_cluster.last_query_cmd = None
+    ota_cluster.emit(
+        OtaQueryCacheClearedEvent.event_type,
+        OtaQueryCacheClearedEvent(
+            device_ieee=str(ieee),
+            endpoint_id=1,
+        ),
+    )
+    await app.shutdown()
+
+    # Reload: the cleared cache should not be restored
+    app2 = await make_app_with_db(db)
+    dev2 = app2.get_device(ieee)
+    ota2 = dev2.endpoints[1].out_clusters[Ota.cluster_id]
+    assert ota2.last_query_cmd is None
+    assert dev2.get_last_ota_query_cmd() is None
+
+    await app2.shutdown()
+
+
+@patch("zigpy.quirks.DEVICE_REGISTRY", new=DeviceRegistry())
+async def test_ota_query_cache_skips_quirk_removed_endpoint(tmp_path):
+    """Test that OTA cache load skips entries for endpoints removed by quirks."""
+    # Register a quirk that removes endpoint 2
+    (
+        QuirkBuilder(
+            "ota manufacturer", "ota model", registry=zigpy.quirks.DEVICE_REGISTRY
+        )
+        .removes_endpoint(2)
+        .add_to_registry()
+    )
+
+    db = tmp_path / "test.db"
+    app = await make_app_with_db(db)
+
+    dev = app.add_device(nwk=0x1234, ieee=t.EUI64.convert("aa:bb:cc:dd:11:22:33:44"))
+    dev.node_desc = make_node_desc(logical_type=zdo_t.LogicalType.Router)
+    dev.model = "ota model"
+    dev.manufacturer = "ota manufacturer"
+
+    ep1 = dev.add_endpoint(1)
+    ep1.status = zigpy.endpoint.Status.ZDO_INIT
+    ep1.profile_id = 260
+    ep1.device_type = profiles.zha.DeviceType.PUMP
+    basic = ep1.add_input_cluster(Basic.cluster_id)
+    basic.update_attribute(Basic.AttributeDefs.model, "ota model")
+    basic.update_attribute(Basic.AttributeDefs.manufacturer, "ota manufacturer")
+
+    ep2 = dev.add_endpoint(2)
+    ep2.status = zigpy.endpoint.Status.ZDO_INIT
+    ep2.profile_id = 260
+    ep2.device_type = profiles.zha.DeviceType.PUMP
+    ota_cluster = ep2.add_output_cluster(Ota.cluster_id)
+
+    ota_cluster.last_query_cmd = Ota.QueryNextImageCommand(
+        field_control=Ota.QueryNextImageCommand.FieldControl(0),
+        manufacturer_code=0x1234,
+        image_type=0x5678,
+        current_file_version=0x00000001,
+    )
+
+    app.device_initialized(dev)
+    await app.shutdown()
+
+    # Reload: quirk removes endpoint 2, OTA cache load should skip it
+    app2 = await make_app_with_db(db)
+    dev2 = app2.get_device(t.EUI64.convert("aa:bb:cc:dd:11:22:33:44"))
+    assert 2 not in dev2.endpoints
+    assert dev2.get_last_ota_query_cmd() is None
+
+    await app2.shutdown()
+
+
+async def test_get_last_ota_query_cmd_returns_none(tmp_path):
+    """Test that get_last_ota_query_cmd returns None when no query has been cached."""
+    db = tmp_path / "test.db"
+    app = await make_app_with_db(db)
+    ieee = make_ieee()
+    app.handle_join(99, ieee, 0)
+
+    dev = app.get_device(ieee)
+    ep = dev.add_endpoint(1)
+    ep.status = zigpy.endpoint.Status.ZDO_INIT
+    ep.profile_id = 260
+    ep.device_type = profiles.zha.DeviceType.PUMP
+    ep.add_output_cluster(Ota.cluster_id)
+    app.device_initialized(dev)
+
+    assert dev.get_last_ota_query_cmd() is None
+
+    await app.shutdown()

{zigpy-1.1.1 → zigpy-1.2.1}/tests/test_appdb_migration.py

@@ -1,6 +1,7 @@
 from datetime import UTC, datetime
 import logging
 import pathlib
+import sqlite3
 from sqlite3.dump import _iterdump as iterdump
 
 from aiosqlite.context import contextmanager
@@ -8,9 +9,8 @@ import pytest
 
 from tests.async_mock import AsyncMock, MagicMock, patch
 from tests.conftest import app, make_node_desc  # noqa: F401
-from tests.test_appdb import auto_kill_aiosqlite, make_app_with_db  # noqa: F401
+from tests.test_appdb import make_app_with_db
 import zigpy.appdb
-from zigpy.appdb import sqlite3
 import zigpy.appdb_schemas
 import zigpy.endpoint
 from zigpy.profiles import zha as zha_profile
@@ -22,6 +22,8 @@ from zigpy.zcl.clusters.general import Basic
 from zigpy.zcl.foundation import BaseAttributeDefs, Status, ZCLAttributeDef
 from zigpy.zdo import types as zdo_t
 
+pytestmark = pytest.mark.usefixtures("auto_kill_aiosqlite")
+
 
 @pytest.fixture
 def test_db(tmp_path):
@@ -276,8 +278,13 @@ async def test_migration_missing_node_descriptor(test_db, caplog):
     [
         ("INSERT INTO node_descriptors_v4 VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)", 0),
         ("INSERT INTO neighbors_v4 VALUES (?,?,?,?,?,?,?,?,?,?,?,?)", 5),
-        ("SELECT * FROM output_clusters", 0),
-        ("INSERT INTO neighbors_v5 VALUES (?,?,?,?,?,?,?,?,?,?,?,?)", 5),
+        ("SELECT ieee, endpoint_id, cluster FROM output_clusters", 0),
+        (
+            "INSERT INTO neighbors_v5 (device_ieee, extended_pan_id, ieee, nwk,"
+            " device_type, rx_on_when_idle, relationship, reserved1, permit_joining,"
+            " reserved2, depth, lqi) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+            5,
+        ),
     ],
 )
 async def test_migration_failure(fail_on_sql, fail_on_count, test_db):
@@ -476,14 +483,36 @@ async def test_migration_missing_tables(app):
     # The untouched table will never be queried
     await appdb._migrate_tables({"table1_v1": "table1_v2", "table2_v1": None})
 
-    mock_execute.assert_called_once_with("SELECT * FROM table1_v1")
-
-    with pytest.raises(AssertionError):
-        mock_execute.assert_called_once_with("SELECT * FROM table2_v1")
+    # table1_v1 is queried (pragma + select), table2_v1 is not
+    assert any("table1_v1" in call.args[0] for call in mock_execute.call_args_list)
+    assert not any("table2_v1" in call.args[0] for call in mock_execute.call_args_list)
 
     await appdb.shutdown()
 
 
+async def test_migrate_tables_integrity_error_raises(tmp_path):
+    """Test that _migrate_tables re-raises IntegrityError with errors='raise'."""
+    db_path = tmp_path / "test.db"
+    app = await make_app_with_db(db_path)
+
+    # Create two simple tables, pre-populate the destination with a conflicting row
+    await app._dblistener.execute(
+        "CREATE TABLE src_v1 (id INTEGER PRIMARY KEY, val TEXT)"
+    )
+    await app._dblistener.execute(
+        "CREATE TABLE dst_v2 (id INTEGER PRIMARY KEY, val TEXT)"
+    )
+    await app._dblistener.execute("INSERT INTO src_v1 VALUES (1, 'hello')")
+    await app._dblistener.execute("INSERT INTO dst_v2 VALUES (1, 'conflict')")
+
+    app._dblistener._get_table_versions = AsyncMock(return_value={"src_v1": "1"})
+
+    with pytest.raises(sqlite3.IntegrityError):
+        await app._dblistener._migrate_tables({"src_v1": "dst_v2"})
+
+    await app.shutdown()
+
+
 async def test_last_seen_initial_migration(test_db):
     test_db_v5 = test_db("simple_v5.sql")
 
@@ -541,7 +570,7 @@ async def test_unknown_manufacturer_code_migration(test_db, caplog):
     # Count rows after migration
     with sqlite3.connect(test_db_prod) as conn:
         cur = conn.cursor()
-        cur.execute("SELECT COUNT(*) FROM attributes_cache_v14")
+        cur.execute(f"SELECT COUNT(*) FROM attributes_cache{zigpy.appdb.DB_V}")
         after_total = cur.fetchone()[0]
 
         assert after_total == before_total
@@ -590,9 +619,9 @@ async def test_manufacturer_code_migration_uses_device_manufacturer_id(test_db):
     with sqlite3.connect(test_db_path) as conn:
         cur = conn.cursor()
         cur.execute(
-            """
+            f"""
             SELECT manufacturer_code
-            FROM attributes_cache_v14
+            FROM attributes_cache{zigpy.appdb.DB_V}
             WHERE cluster_id = 0xFC00 AND attr_id = 0x0002
             """,
         )
@@ -604,9 +633,9 @@ async def test_manufacturer_code_migration_uses_device_manufacturer_id(test_db):
     with sqlite3.connect(test_db_path) as conn:
         cur = conn.cursor()
         cur.execute(
-            """
+            f"""
             SELECT manufacturer_code, status
-            FROM attributes_cache_v14
+            FROM attributes_cache{zigpy.appdb.DB_V}
             WHERE ieee = ? AND cluster_id = 0xFC00 AND attr_id = 0x0004
             """,
             (third_reality_ieee,),
@@ -697,7 +726,7 @@ async def test_data_migration_ambiguous_attributes(tmp_path):
 
     with sqlite3.connect(db_path) as conn:
         conn.executemany(
-            "INSERT INTO attributes_cache_v14"
+            f"INSERT INTO attributes_cache{zigpy.appdb.DB_V}"
             " (ieee, endpoint_id, cluster_type, cluster_id,"
             "  attr_id, manufacturer_code, status, value, last_updated)"
             " VALUES (?, ?, ?, ?, ?, ?, 0, ?, ?)",
@@ -736,7 +765,7 @@ async def test_data_migration_ambiguous_attributes(tmp_path):
     with sqlite3.connect(db_path) as conn:
         # The disambiguated unmigrated row was deleted (a row with 0xABCD already existed)
         rows = conn.execute(
-            "SELECT manufacturer_code FROM attributes_cache_v14"
+            f"SELECT manufacturer_code FROM attributes_cache{zigpy.appdb.DB_V}"
             " WHERE ieee = ? AND cluster_id = ? AND attr_id = ?",
             (str(dev.ieee), 0xFC01, 0x0010),
         ).fetchall()
@@ -744,7 +773,7 @@ async def test_data_migration_ambiguous_attributes(tmp_path):
 
         # The ambiguous unmigrated row is still present
         rows = conn.execute(
-            "SELECT manufacturer_code FROM attributes_cache_v14"
+            f"SELECT manufacturer_code FROM attributes_cache{zigpy.appdb.DB_V}"
             " WHERE ieee = ? AND cluster_id = ? AND attr_id = ?",
             (str(dev.ieee), 0xFC02, 0x0020),
         ).fetchall()

{zigpy-1.1.1 → zigpy-1.2.1}/tests/test_device.py

@@ -18,7 +18,7 @@ from zigpy.profiles import zha
 import zigpy.state
 import zigpy.types as t
 import zigpy.util
-from zigpy.zcl import ClusterType, foundation
+from zigpy.zcl import ClusterType, OtaQueryCacheClearedEvent, foundation
 from zigpy.zcl.clusters.general import Basic, OnOff, Ota, PollControl
 from zigpy.zdo import types as zdo_t
 
@@ -108,6 +108,57 @@ async def test_initialize_read_ota(
     assert success[Ota.AttributeDefs.current_file_version.id] == 0x12345678
 
 
+async def test_initialize_sends_image_notify(
+    app: zigpy.application.ControllerApplication,
+) -> None:
+    """Test that image_notify is sent to OTA devices during initialization."""
+    dev = app.add_device(nwk=0x1234, ieee=t.EUI64.convert("aa:bb:cc:dd:ee:ff:00:11"))
+    dev.node_desc = make_node_desc()
+
+    ep = dev.add_endpoint(1)
+    ep.status = endpoint.Status.ZDO_INIT
+
+    basic = ep.add_input_cluster(Basic.cluster_id)
+    ota = ep.add_output_cluster(Ota.cluster_id)
+
+    with (
+        mock_attribute_reads(basic, {"model": "Model", "manufacturer": "Manufacturer"}),
+        mock_attribute_reads(ota, {"current_file_version": 0x12345678}),
+    ):
+        await dev.initialize()
+
+    # image_notify is the last packet sent during initialization
+    packet = app.send_packet.call_args_list[0][0][0]
+    hdr, cmd = ota.deserialize(packet.data.serialize())
+    assert isinstance(cmd, Ota.ImageNotifyCommand)
+    assert cmd.payload_type == Ota.ImageNotifyCommand.PayloadType.QueryJitter
+    assert cmd.query_jitter == 100
+
+
+async def test_initialize_image_notify_failure_does_not_block(
+    app: zigpy.application.ControllerApplication,
+) -> None:
+    """Test that a failed image_notify does not prevent initialization."""
+    dev = app.add_device(nwk=0x1234, ieee=t.EUI64.convert("aa:bb:cc:dd:ee:ff:00:11"))
+    dev.node_desc = make_node_desc()
+
+    ep = dev.add_endpoint(1)
+    ep.status = endpoint.Status.ZDO_INIT
+
+    basic = ep.add_input_cluster(Basic.cluster_id)
+    ota = ep.add_output_cluster(Ota.cluster_id)
+    ota.image_notify = AsyncMock(side_effect=TimeoutError)
+
+    with (
+        mock_attribute_reads(basic, {"model": "Model", "manufacturer": "Manufacturer"}),
+        mock_attribute_reads(ota, {"current_file_version": 0x12345678}),
+    ):
+        await dev.initialize()
+
+    assert dev.is_initialized
+    ota.image_notify.assert_awaited_once()
+
+
 async def test_initialize_read_ota_unsupported(
     app: zigpy.application.ControllerApplication,
 ) -> None:
@@ -452,19 +503,21 @@ async def test_handle_unknown_cluster(dev, caplog) -> None:
 
 async def test_update_device_firmware_no_ota_cluster(dev):
     """Test that device firmware updates fails: no ota cluster."""
+    mock_image = MagicMock()
+
     with pytest.raises(ValueError, match="Cluster 0x0019 not found"):
-        await dev.update_firmware(sentinel.firmware_image, sentinel.progress_callback)
+        await dev.update_firmware(mock_image, sentinel.progress_callback)
 
     dev.add_endpoint(1)
     dev.endpoints[1].output_clusters = MagicMock(side_effect=KeyError)
     with pytest.raises(ValueError, match="Cluster 0x0019 not found"):
-        await dev.update_firmware(sentinel.firmware_image, sentinel.progress_callback)
+        await dev.update_firmware(mock_image, sentinel.progress_callback)
 
 
 async def test_update_device_firmware_already_in_progress(dev, caplog):
     """Test that device firmware updates no ops when update is in progress."""
     dev.ota_in_progress = True
-    await dev.update_firmware(sentinel.firmware_image, sentinel.progress_callback)
+    await dev.update_firmware(MagicMock(), sentinel.progress_callback)
     assert "OTA already in progress" in caplog.text
 
 
@@ -529,7 +582,7 @@ async def test_update_device_firmware(monkeypatch, dev, caplog):
         )
     )
 
-    dev.application.ota.get_ota_images = MagicMock(
+    dev.application.ota.get_ota_images = AsyncMock(
         return_value=OtaImagesResult(upgrades=(), downgrades=())
     )
     dev.update_firmware = MagicMock(wraps=dev.update_firmware)
@@ -582,6 +635,9 @@ async def test_update_device_firmware(monkeypatch, dev, caplog):
             elif isinstance(
                 cmd, Ota.ClientCommandDefs.query_next_image_response.schema
             ):
+                # Post-OTA image_notify triggers a query with NO_IMAGE_AVAILABLE
+                if cmd.status == foundation.Status.NO_IMAGE_AVAILABLE:
+                    return
                 assert cmd.status == foundation.Status.SUCCESS
                 assert (
                     cmd.manufacturer_code
@@ -711,6 +767,10 @@ async def test_update_device_firmware(monkeypatch, dev, caplog):
 
     dev.application.send_packet = AsyncMock(side_effect=send_packet)
     progress_callback = MagicMock()
+
+    cleared_events = []
+    cluster.on_event(OtaQueryCacheClearedEvent.event_type, cleared_events.append)
+
     result = await dev.update_firmware(fw_image, progress_callback)
     assert (
         dev.endpoints[1]
@@ -719,11 +779,18 @@ async def test_update_device_firmware(monkeypatch, dev, caplog):
         == 0x12345678
     )
 
-    assert dev.application.send_packet.await_count == 6
+    # Wait for background tasks (post-OTA query_next_image handling)
+    await asyncio.sleep(0)
+    # 6 OTA + 1 post-OTA image_notify + 1 post-OTA query_next_image_response
+    assert dev.application.send_packet.await_count == 8
     assert progress_callback.call_count == 2
     assert progress_callback.call_args_list[0] == call(40, 70, 57.142857142857146)
     assert progress_callback.call_args_list[1] == call(70, 70, 100.0)
     assert result == foundation.Status.SUCCESS
+    assert len(cleared_events) == 1
+    assert isinstance(cleared_events[0], OtaQueryCacheClearedEvent)
+    # Post-OTA image_notify repopulated the query cache
+    assert cluster.last_query_cmd is not None
 
     progress_callback.reset_mock()
     dev.application.send_packet.reset_mock()
@@ -731,12 +798,35 @@ async def test_update_device_firmware(monkeypatch, dev, caplog):
         fw_image, progress_callback=progress_callback, force=True
     )
 
-    assert dev.application.send_packet.await_count == 6
+    await asyncio.sleep(0)
+    # 6 OTA + 1 post-OTA image_notify + 1 post-OTA query_next_image_response
+    assert dev.application.send_packet.await_count == 8
     assert progress_callback.call_count == 2
     assert progress_callback.call_args_list[0] == call(40, 70, 57.142857142857146)
     assert progress_callback.call_args_list[1] == call(70, 70, 100.0)
     assert result == foundation.Status.SUCCESS
 
+    # Post-OTA image_notify failure: OTA still succeeds
+    dev.application.send_packet.reset_mock()
+    progress_callback.reset_mock()
+    caplog.clear()
+    original_image_notify = cluster.image_notify
+    notify_calls = 0
+
+    async def image_notify_fail_post_ota(*args, **kwargs):
+        nonlocal notify_calls
+        notify_calls += 1
+        if notify_calls > 1:
+            raise zigpy.exceptions.DeliveryError("Device rebooting")
+        return await original_image_notify(*args, **kwargs)
+
+    cluster.image_notify = image_notify_fail_post_ota
+    result = await dev.update_firmware(fw_image, progress_callback=progress_callback)
+    assert result == foundation.Status.SUCCESS
+    assert "Post-OTA image_notify failed" in caplog.text
+    cluster.image_notify = original_image_notify
+    caplog.clear()
+
     # _image_query_req exception test
     dev.application.send_packet.reset_mock()
     progress_callback.reset_mock()
@@ -816,6 +906,9 @@ async def test_update_device_firmware(monkeypatch, dev, caplog):
             elif isinstance(
                 cmd, Ota.ClientCommandDefs.query_next_image_response.schema
             ):
+                # Post-OTA image_notify triggers a query with NO_IMAGE_AVAILABLE
+                if cmd.status == foundation.Status.NO_IMAGE_AVAILABLE:
+                    return
                 assert cmd.status == foundation.Status.SUCCESS
                 assert cmd.manufacturer_code == fw_image.firmware.header.manufacturer_id
                 assert cmd.image_type == fw_image.firmware.header.image_type
@@ -912,7 +1005,7 @@ async def test_update_legrand_device_firmware(monkeypatch, dev, caplog):
         )
     )
 
-    dev.application.ota.get_ota_images = MagicMock(
+    dev.application.ota.get_ota_images = AsyncMock(
         return_value=OtaImagesResult(upgrades=(), downgrades=())
     )
     dev.update_firmware = MagicMock(wraps=dev.update_firmware)
@@ -964,6 +1057,9 @@ async def test_update_legrand_device_firmware(monkeypatch, dev, caplog):
             elif isinstance(
                 cmd, Ota.ClientCommandDefs.query_next_image_response.schema
             ):
+                # Post-OTA image_notify triggers a query with NO_IMAGE_AVAILABLE
+                if cmd.status == foundation.Status.NO_IMAGE_AVAILABLE:
+                    return
                 assert cmd.status == foundation.Status.SUCCESS
                 assert (
                     cmd.manufacturer_code
@@ -1101,7 +1197,10 @@ async def test_update_legrand_device_firmware(monkeypatch, dev, caplog):
         == 0x12345678
     )
 
-    assert dev.application.send_packet.await_count == 6
+    # Wait for background tasks (post-OTA query_next_image handling)
+    await asyncio.sleep(0)
+    # 6 OTA + 1 post-OTA image_notify + 1 post-OTA query_next_image_response
+    assert dev.application.send_packet.await_count == 8
     assert progress_callback.call_count == 2
     assert progress_callback.call_args_list[0] == call(64, 70, 91.42857142857143)
     assert progress_callback.call_args_list[1] == call(70, 70, 100.0)
@@ -1113,7 +1212,9 @@ async def test_update_legrand_device_firmware(monkeypatch, dev, caplog):
         fw_image, progress_callback=progress_callback, force=True
     )
 
-    assert dev.application.send_packet.await_count == 6
+    await asyncio.sleep(0)
+    # 6 OTA + 1 post-OTA image_notify + 1 post-OTA query_next_image_response
+    assert dev.application.send_packet.await_count == 8
     assert progress_callback.call_count == 2
     assert progress_callback.call_args_list[0] == call(64, 70, 91.42857142857143)
     assert progress_callback.call_args_list[1] == call(70, 70, 100.0)
@@ -1198,6 +1299,9 @@ async def test_update_legrand_device_firmware(monkeypatch, dev, caplog):
             elif isinstance(
                 cmd, Ota.ClientCommandDefs.query_next_image_response.schema
             ):
+                # Post-OTA image_notify triggers a query with NO_IMAGE_AVAILABLE
+                if cmd.status == foundation.Status.NO_IMAGE_AVAILABLE:
+                    return
                 assert cmd.status == foundation.Status.SUCCESS
                 assert cmd.manufacturer_code == fw_image.firmware.header.manufacturer_id
                 assert cmd.image_type == fw_image.firmware.header.image_type