minor change to frag ack callback handling, added test case for fragmentation

Author: tkulin

bellows/ezsp/protocol.py

@@ -212,10 +212,13 @@ def __call__(self, data: bytes) -> None:
                 fragment_index=fragment_index,
                 payload=result[7],
             )
+            if not hasattr(self, "_ack_tasks"):
+                self._ack_tasks = set()
             ack_task = asyncio.create_task(
                 self._send_fragment_ack(sender, aps_frame, frag_count, frag_index)
             )  # APS Ack
-            ack_task.add_done_callback(self._ack_tasks.remove)
+            self._ack_tasks.add(ack_task)
+            ack_task.add_done_callback(lambda t: self._ack_tasks.discard(t))
 
             if not complete:
                 # Do not pass partial data up the stack

tests/test_fragmentation.py

@@ -0,0 +1,192 @@
+import pytest
+
+from bellows.ezsp.fragmentation import FragmentManager
+
+
+@pytest.fixture
+def frag_manager():
+    return FragmentManager()
+
+
+@pytest.mark.asyncio
+async def test_single_fragment_complete(frag_manager):
+    # If we receive a single-fragment message, the fragemnt manager should immediately report completion.
+
+    key = (0x1234, 0xAB, 0x1234, 0x5678)
+    fragment_count = 1
+    fragment_index = 0
+    payload = b"Single fragment"
+
+    (
+        complete,
+        reassembled,
+        returned_frag_count,
+        returned_frag_index,
+    ) = frag_manager.handle_incoming_fragment(
+        sender_nwk=key[0],
+        aps_sequence=key[1],
+        profile_id=key[2],
+        cluster_id=key[3],
+        fragment_count=fragment_count,
+        fragment_index=fragment_index,
+        payload=payload,
+    )
+
+    assert complete is True
+    assert reassembled == payload
+    assert returned_frag_count == fragment_count
+    assert returned_frag_index == fragment_index
+    assert key not in frag_manager._partial
+    assert key not in frag_manager._cleanup_timers
+
+
+@pytest.mark.asyncio
+async def test_two_fragments_in_order(frag_manager):
+    # A two-fragment message should remain partial until we've received both pieces.
+    key = (0x1111, 0x01, 0x9999, 0x2222)
+    fragment_count = 2
+
+    # First fragment
+    (
+        complete,
+        reassembled,
+        returned_frag_count,
+        returned_frag_index,
+    ) = frag_manager.handle_incoming_fragment(
+        sender_nwk=key[0],
+        aps_sequence=key[1],
+        profile_id=key[2],
+        cluster_id=key[3],
+        fragment_count=fragment_count,
+        fragment_index=0,
+        payload=b"Frag0-",
+    )
+    assert complete is False
+    assert reassembled is None
+    assert key in frag_manager._partial
+    assert frag_manager._partial[key].fragments_received == 1
+
+    # Second fragment
+    (
+        complete,
+        reassembled,
+        returned_frag_count,
+        returned_frag_index,
+    ) = frag_manager.handle_incoming_fragment(
+        sender_nwk=key[0],
+        aps_sequence=key[1],
+        profile_id=key[2],
+        cluster_id=key[3],
+        fragment_count=fragment_count,
+        fragment_index=1,
+        payload=b"Frag1",
+    )
+    assert complete is True
+    assert reassembled == b"Frag0-Frag1"
+    assert key not in frag_manager._partial
+    assert key not in frag_manager._cleanup_timers
+
+
+@pytest.mark.asyncio
+async def test_out_of_order_fragments(frag_manager):
+    # Receiving fragments in reverse order
+    key = (0x9999, 0xCD, 0x1234, 0xABCD)
+    fragment_count = 2
+
+    # Second fragment arrives first
+    (
+        complete,
+        reassembled,
+        returned_frag_count,
+        returned_frag_index,
+    ) = frag_manager.handle_incoming_fragment(
+        sender_nwk=key[0],
+        aps_sequence=key[1],
+        profile_id=key[2],
+        cluster_id=key[3],
+        fragment_count=fragment_count,
+        fragment_index=1,
+        payload=b"World",
+    )
+    assert not complete
+    assert reassembled is None
+
+    # Then the first fragment
+    (
+        complete,
+        reassembled,
+        returned_frag_count,
+        returned_frag_index,
+    ) = frag_manager.handle_incoming_fragment(
+        sender_nwk=key[0],
+        aps_sequence=key[1],
+        profile_id=key[2],
+        cluster_id=key[3],
+        fragment_count=fragment_count,
+        fragment_index=0,
+        payload=b"Hello ",
+    )
+    assert complete
+    assert reassembled == b"Hello World"
+
+
+@pytest.mark.asyncio
+async def test_repeated_fragments_ignored(frag_manager):
+    # Ensure repeated arrivals of the same fragment index do not double-count or break the logic.
+
+    key = (0xAAA, 0xBB, 0xCCC, 0xDDD)
+    fragment_count = 2
+
+    # First fragment
+    (
+        complete,
+        reassembled,
+        returned_frag_count,
+        returned_frag_index,
+    ) = frag_manager.handle_incoming_fragment(
+        sender_nwk=key[0],
+        aps_sequence=key[1],
+        profile_id=key[2],
+        cluster_id=key[3],
+        fragment_count=fragment_count,
+        fragment_index=0,
+        payload=b"first",
+    )
+    assert not complete
+    assert frag_manager._partial[key].fragments_received == 1
+
+    # Repeat the same fragment index
+    (
+        complete,
+        reassembled,
+        returned_frag_count,
+        returned_frag_index,
+    ) = frag_manager.handle_incoming_fragment(
+        sender_nwk=key[0],
+        aps_sequence=key[1],
+        profile_id=key[2],
+        cluster_id=key[3],
+        fragment_count=fragment_count,
+        fragment_index=0,
+        payload=b"first",
+    )
+    assert not complete
+    assert frag_manager._partial[key].fragments_received == 1, "Should not increment"
+
+    # Second fragment completes
+    (
+        complete,
+        reassembled,
+        returned_frag_count,
+        returned_frag_index,
+    ) = frag_manager.handle_incoming_fragment(
+        sender_nwk=key[0],
+        aps_sequence=key[1],
+        profile_id=key[2],
+        cluster_id=key[3],
+        fragment_count=fragment_count,
+        fragment_index=1,
+        payload=b"second",
+    )
+    assert complete
+    assert reassembled == b"firstsecond"