.

Author: hampsterx

kinesis/consumer.py

@@ -147,9 +147,10 @@ async def close(self):
 
             # Flush deferred checkpoints (from __anext__ deferral)
             if self._deferred_checkpoints and self.checkpointer:
-                for shard_id, sequence in self._deferred_checkpoints.items():
+                for shard_id, sequence in list(self._deferred_checkpoints.items()):
                     await self._maybe_checkpoint(shard_id, sequence)
-                self._deferred_checkpoints.clear()
+                    if self._deferred_checkpoints.get(shard_id) == sequence:
+                        del self._deferred_checkpoints[shard_id]
 
             # Cancel background flusher (triggers final flush via CancelledError)
             if self._checkpoint_flusher_task is not None:
@@ -373,15 +374,17 @@ async def fetch(self):
                         if self.checkpointer:
                             # Flush deferred checkpoint if it's for this shard
                             if shard_id in self._deferred_checkpoints:
-                                await self._maybe_checkpoint(
-                                    shard_id, self._deferred_checkpoints.pop(shard_id)
-                                )
+                                seq = self._deferred_checkpoints[shard_id]
+                                await self._maybe_checkpoint(shard_id, seq)
+                                if self._deferred_checkpoints.get(shard_id) == seq:
+                                    del self._deferred_checkpoints[shard_id]
 
                             # Flush interval-buffered checkpoint for this shard
                             if shard_id in self._pending_checkpoints:
-                                await self.checkpointer.checkpoint(
-                                    shard_id, self._pending_checkpoints.pop(shard_id)
-                                )
+                                seq = self._pending_checkpoints[shard_id]
+                                await self.checkpointer.checkpoint(shard_id, seq)
+                                if self._pending_checkpoints.get(shard_id) == seq:
+                                    del self._pending_checkpoints[shard_id]
 
                             # Note: the terminal batch's records are enqueued but no
                             # checkpoint sentinel was added (would race with deallocate).
@@ -877,7 +880,9 @@ async def _flush_pending_checkpoints(self):
         newer sequence written by _maybe_checkpoint during the await.
         """
         for shard_id in list(self._pending_checkpoints):
-            seq = self._pending_checkpoints[shard_id]
+            seq = self._pending_checkpoints.get(shard_id)
+            if seq is None:
+                continue
             await self.checkpointer.checkpoint(shard_id, seq)
             # Only delete if the value hasn't been superseded during the await
             if self._pending_checkpoints.get(shard_id) == seq:
@@ -900,9 +905,10 @@ async def __anext__(self):
 
         # 1. Commit deferred checkpoints from previous iteration (implicit ACK)
         if self._deferred_checkpoints:
-            for shard_id, sequence in self._deferred_checkpoints.items():
+            for shard_id, sequence in list(self._deferred_checkpoints.items()):
                 await self._maybe_checkpoint(shard_id, sequence)
-            self._deferred_checkpoints.clear()
+                if self._deferred_checkpoints.get(shard_id) == sequence:
+                    del self._deferred_checkpoints[shard_id]
 
         checkpoint_count = 0
         max_checkpoints = 100  # Prevent infinite checkpoint processing

tests/conftest.py

@@ -189,12 +189,15 @@ def _factory(**kwargs):
 
     yield _factory
 
+    teardown_errors = []
     for c in consumers:
         for task in [getattr(c, "_checkpoint_flusher_task", None), c.fetch_task]:
             if task is None:
                 continue
             if task.done():
-                # Surface exceptions from tasks that finished during the test
+                # Log exceptions from tasks that finished during the test.
+                # Don't re-raise: the test itself is responsible for observing
+                # expected task failures (e.g. test_wait_ready_fetch_task_crash).
                 if not task.cancelled() and task.exception():
                     log.warning(
                         "Task finished with error during test: %s", task.exception(), exc_info=task.exception()
@@ -205,5 +208,8 @@ def _factory(**kwargs):
                     await task
                 except asyncio.CancelledError:
                     pass
-                except Exception:
-                    log.warning("Unexpected error during mock_consumer teardown", exc_info=True)
+                except Exception as exc:
+                    teardown_errors.append(exc)
+    if teardown_errors:
+        log.warning("mock_consumer teardown errors: %s", teardown_errors)
+        raise teardown_errors[0]