From 69de53637ce021049de45a0684bfae6baa57bfe0 Mon Sep 17 00:00:00 2001
From: Kaushik Raina <kraina@confluent.io>
Date: Mon, 20 Apr 2026 10:35:58 +0530
Subject: [PATCH 1/6] Fix consume() breaking out of chunked loop on first
 result

---
 src/confluent_kafka/src/Consumer.c            |  42 ++++--
 .../test_consumer_wakeable_poll_consume.py    | 138 ++++++++++++++++++
 tests/test_Wakeable.py                        | 126 ++++++++++++++++
 3 files changed, 291 insertions(+), 15 deletions(-)

diff --git a/src/confluent_kafka/src/Consumer.c b/src/confluent_kafka/src/Consumer.c
index d71af1243..f83fbcf9a 100644
--- a/src/confluent_kafka/src/Consumer.c
+++ b/src/confluent_kafka/src/Consumer.c
@@ -1119,7 +1119,8 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
         PyObject *msglist;
         rd_kafka_queue_t *rkqu = self->u.Consumer.rkqu;
         CallState cs;
-        Py_ssize_t i, n = 0;
+        Py_ssize_t i, msgs_received_count = 0;
+        Py_ssize_t chunk_msg_count;
         const int CHUNK_TIMEOUT_MS = 200; /* 200ms chunks for signal checking */
         int total_timeout_ms;
         int chunk_timeout_ms;
@@ -1156,10 +1157,10 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
          * ThreadPool. Only use wakeable poll for
          * blocking calls that need to be interruptible. */
         if (total_timeout_ms >= 0 && total_timeout_ms < CHUNK_TIMEOUT_MS) {
-                n = (Py_ssize_t)rd_kafka_consume_batch_queue(
+                msgs_received_count = (Py_ssize_t)rd_kafka_consume_batch_queue(
                     rkqu, total_timeout_ms, rkmessages, num_messages);
 
-                if (n < 0) {
+                if (msgs_received_count < 0) {
                         /* Error - need to restore GIL before setting error */
                         PyEval_RestoreThread(cs.thread_state);
                         free(rkmessages);
@@ -1178,13 +1179,20 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
                                 break;
                         }
 
-                        /* Consume with chunk timeout */
-                        n = (Py_ssize_t)rd_kafka_consume_batch_queue(
-                            rkqu, chunk_timeout_ms, rkmessages, num_messages);
-
-                        if (n < 0) {
-                                /* Error - need to restore GIL before setting
-                                 * error */
+                        /* Consume with chunk timeout, appending after
+                         * already-accumulated messages */
+                        chunk_msg_count =
+                            (Py_ssize_t)rd_kafka_consume_batch_queue(
+                                rkqu, chunk_timeout_ms,
+                                rkmessages + msgs_received_count,
+                                num_messages -
+                                    (unsigned int)msgs_received_count);
+
+                        if (chunk_msg_count < 0) {
+                                /* Error - destroy accumulated messages,
+                                 * restore GIL, and raise */
+                                for (i = 0; i < msgs_received_count; i++)
+                                        rd_kafka_message_destroy(rkmessages[i]);
                                 PyEval_RestoreThread(cs.thread_state);
                                 free(rkmessages);
                                 cfl_PyErr_Format(
@@ -1193,8 +1201,10 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
                                 return NULL;
                         }
 
-                        /* If we got messages, exit the loop */
-                        if (n > 0) {
+                        msgs_received_count += chunk_msg_count;
+
+                        /* If we got all requested messages, exit the loop */
+                        if (msgs_received_count >= (Py_ssize_t)num_messages) {
                                 break;
                         }
 
@@ -1202,6 +1212,8 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
 
                         /* Check for signals between chunks */
                         if (check_signals_between_chunks(self, &cs)) {
+                                for (i = 0; i < msgs_received_count; i++)
+                                        rd_kafka_message_destroy(rkmessages[i]);
                                 free(rkmessages);
                                 return NULL;
                         }
@@ -1210,7 +1222,7 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
 
         /* Final GIL restore and signal check */
         if (!CallState_end(self, &cs)) {
-                for (i = 0; i < n; i++) {
+                for (i = 0; i < msgs_received_count; i++) {
                         rd_kafka_message_destroy(rkmessages[i]);
                 }
                 free(rkmessages);
@@ -1218,9 +1230,9 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
         }
 
         /* Create Python list from messages  */
-        msglist = PyList_New(n);
+        msglist = PyList_New(msgs_received_count);
 
-        for (i = 0; i < n; i++) {
+        for (i = 0; i < msgs_received_count; i++) {
                 PyObject *msgobj = Message_new0(self, rkmessages[i]);
 #ifdef RD_KAFKA_V_HEADERS
                 /** Have to detach headers outside Message_new0 because it
diff --git a/tests/integration/consumer/test_consumer_wakeable_poll_consume.py b/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
index 8ca5aea86..ab49627b1 100644
--- a/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
+++ b/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
@@ -142,3 +142,141 @@ def test_consume_message_delivery_with_wakeable_pattern(kafka_cluster):
         assert msg.value() == expected_value, expected_msg
 
     consumer.close()
+
+
+def test_consume_accumulates_messages_across_chunks(kafka_cluster):
+    """Test that consume() accumulates messages across 200ms chunks.
+    This verifies that consume() doesn't return early on the first chunk of messages
+    when using the wakeable pattern.
+    """
+    topic = kafka_cluster.create_topic_and_wait_propogation('test-consume-accumulate-chunks')
+
+    # Produce 10 messages
+    producer = kafka_cluster.cimpl_producer()
+    num_produced = 10
+    for i in range(num_produced):
+        producer.produce(topic, value=f'msg-{i}'.encode())
+    producer.flush(timeout=5.0)
+
+    # Create consumer
+    consumer_conf = kafka_cluster.client_conf(
+        {
+            'group.id': 'test-consume-accumulate',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 6000,
+            'auto.offset.reset': 'earliest',
+        }
+    )
+    consumer = TestConsumer(consumer_conf)
+    consumer.subscribe([topic])
+
+    # Wait for subscription and partition assignment
+    time.sleep(2.0)
+
+    # Consume with num_messages=10 and a generous timeout.
+    # Before the fix: would return < 10 (whatever arrived in the first 200ms chunk)
+    # After the fix: accumulates across chunks until 10 are collected
+    msglist = consumer.consume(num_messages=num_produced, timeout=10.0)
+
+    assert len(msglist) == num_produced, (
+        f"Expected {num_produced} messages but got {len(msglist)}. " f"consume() may not be accumulating across chunks."
+    )
+
+    for i, msg in enumerate(msglist):
+        assert not msg.error(), f"Message {i} has error: {msg.error()}"
+
+    consumer.close()
+
+
+def test_consume_returns_partial_on_timeout(kafka_cluster):
+    """Test that consume() returns partial results when timeout expires
+    before num_messages is reached."""
+    topic = kafka_cluster.create_topic_and_wait_propogation('test-consume-partial-timeout')
+
+    # Produce only 3 messages, but request 100
+    producer = kafka_cluster.cimpl_producer()
+    num_produced = 3
+    for i in range(num_produced):
+        producer.produce(topic, value=f'partial-{i}'.encode())
+    producer.flush(timeout=5.0)
+
+    consumer_conf = kafka_cluster.client_conf(
+        {
+            'group.id': 'test-consume-partial',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 6000,
+            'auto.offset.reset': 'earliest',
+        }
+    )
+    consumer = TestConsumer(consumer_conf)
+    consumer.subscribe([topic])
+
+    time.sleep(2.0)
+
+    # Request 100 messages but only 3 exist — should return 3 after timeout
+    start = time.time()
+    msglist = consumer.consume(num_messages=100, timeout=3.0)
+    elapsed = time.time() - start
+
+    assert len(msglist) == num_produced, f"Expected {num_produced} messages (partial), got {len(msglist)}"
+    # Should have waited close to the full timeout since num_messages wasn't reached
+    assert elapsed >= 2.0, f"Should wait near full timeout for more messages, but returned in {elapsed:.2f}s"
+
+    for i, msg in enumerate(msglist):
+        assert not msg.error(), f"Message {i} has error: {msg.error()}"
+        assert msg.value() == f'partial-{i}'.encode()
+
+    consumer.close()
+
+
+def test_consume_accumulates_messages_produced_in_waves(kafka_cluster):
+    """Test that consume() accumulates messages that arrive in multiple waves.
+    This verifies that consume() doesn't return early on the first wave of messages
+    when using the wakeable pattern.
+    """
+    import threading
+
+    topic = kafka_cluster.create_topic_and_wait_propogation('test-consume-waves')
+
+    producer = kafka_cluster.cimpl_producer()
+
+    def produce_in_waves():
+        """Produce 3 waves of 4 messages each, with 1s gaps."""
+        for wave in range(3):
+            time.sleep(1.0)
+            for i in range(4):
+                msg_num = wave * 4 + i
+                producer.produce(topic, value=f'wave-{msg_num}'.encode())
+            producer.flush(timeout=5.0)
+
+    consumer_conf = kafka_cluster.client_conf(
+        {
+            'group.id': 'test-consume-waves',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 6000,
+            'auto.offset.reset': 'earliest',
+        }
+    )
+    consumer = TestConsumer(consumer_conf)
+    consumer.subscribe([topic])
+
+    time.sleep(2.0)
+
+    # Start producing in background
+    producer_thread = threading.Thread(target=produce_in_waves, daemon=True)
+    producer_thread.start()
+
+    # Request 10 messages with a long timeout (waves take ~3s to complete)
+    msglist = consumer.consume(num_messages=10, timeout=10.0)
+
+    # Should have accumulated messages across multiple waves
+    assert len(msglist) >= 10, (
+        f"Expected at least 10 messages accumulated across waves, got {len(msglist)}. "
+        f"consume() may be returning early on the first wave."
+    )
+
+    for msg in msglist:
+        assert not msg.error(), f"Message has error: {msg.error()}"
+
+    producer_thread.join(timeout=5.0)
+    consumer.close()
diff --git a/tests/test_Wakeable.py b/tests/test_Wakeable.py
index ea0947976..3f962f38d 100644
--- a/tests/test_Wakeable.py
+++ b/tests/test_Wakeable.py
@@ -1438,3 +1438,129 @@ def test_flush_empty_queue_returns_immediately():
     # Key assertion: empty flush is fast
     assert qlen == 0, "Empty queue should return 0"
     assert elapsed < WAKEABLE_POLL_TIMEOUT_MAX, f"Empty flush should return quickly, took {elapsed:.2f}s"
+
+
+# These tests verify that consume() correctly accumulates messages across
+# 200ms chunks instead of returning on the first non-empty chunk.
+
+
+def test_consume_accumulates_across_chunks_no_messages():
+    """consume() with no messages available should wait the full timeout
+    and return an empty list, not return early on the first empty chunk."""
+    consumer = TestConsumer(
+        {
+            'group.id': 'test-accumulate-no-msgs',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 1000,
+            'auto.offset.reset': 'latest',
+        }
+    )
+    consumer.subscribe(['test-accumulate-topic'])
+
+    start = time.time()
+    msglist = consumer.consume(num_messages=10, timeout=0.5)
+    elapsed = time.time() - start
+
+    assert isinstance(msglist, list), "consume() should return a list"
+    assert len(msglist) == 0, "Expected empty list when no messages available"
+    assert (
+        WAKEABLE_POLL_TIMEOUT_MIN <= elapsed <= WAKEABLE_POLL_TIMEOUT_MAX
+    ), f"Should wait ~0.5s for timeout, took {elapsed:.2f}s"
+    consumer.close()
+
+
+def test_consume_accumulation_signal_interrupts_cleanly():
+    """When Ctrl+C arrives during accumulation, consume() should raise
+    KeyboardInterrupt without leaking messages."""
+    consumer = TestConsumer(
+        {
+            'group.id': 'test-accumulate-signal',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 1000,
+            'auto.offset.reset': 'latest',
+        }
+    )
+    consumer.subscribe(['test-accumulate-topic'])
+
+    interrupt_thread = threading.Thread(target=lambda: TestUtils.send_sigint_after_delay(0.3))
+    interrupt_thread.daemon = True
+    interrupt_thread.start()
+
+    interrupted = False
+    try:
+        consumer.consume(num_messages=100, timeout=WAKEABLE_POLL_TIMEOUT_MAX)
+    except KeyboardInterrupt:
+        interrupted = True
+    finally:
+        consumer.close()
+
+    assert interrupted, "Should have raised KeyboardInterrupt"
+
+
+def test_consume_accumulation_timeout_returns_partial():
+    """consume() should return whatever messages it has when timeout expires,
+    even if fewer than num_messages. With no broker, this means empty list."""
+    consumer = TestConsumer(
+        {
+            'group.id': 'test-accumulate-partial',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 1000,
+            'auto.offset.reset': 'latest',
+        }
+    )
+    consumer.subscribe(['test-accumulate-topic'])
+
+    start = time.time()
+    msglist = consumer.consume(num_messages=1000, timeout=0.5)
+    elapsed = time.time() - start
+
+    assert isinstance(msglist, list), "consume() should return a list"
+    assert len(msglist) == 0, "Expected empty list with no broker"
+    assert (
+        WAKEABLE_POLL_TIMEOUT_MIN <= elapsed <= WAKEABLE_POLL_TIMEOUT_MAX
+    ), f"Should wait full timeout, took {elapsed:.2f}s"
+    consumer.close()
+
+
+def test_consume_accumulation_short_timeout_skips_chunking():
+    """consume() with timeout < 200ms should skip the chunked loop entirely."""
+    consumer = TestConsumer(
+        {
+            'group.id': 'test-accumulate-short',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 1000,
+            'auto.offset.reset': 'latest',
+        }
+    )
+    consumer.subscribe(['test-accumulate-topic'])
+
+    start = time.time()
+    msglist = consumer.consume(num_messages=10, timeout=0.05)
+    elapsed = time.time() - start
+
+    assert isinstance(msglist, list), "consume() should return a list"
+    assert len(msglist) == 0, "Expected empty list"
+    assert elapsed < WAKEABLE_POLL_TIMEOUT_MIN, f"Short timeout should not use chunking, took {elapsed:.2f}s"
+    consumer.close()
+
+
+def test_consume_accumulation_zero_timeout_nonblocking():
+    """consume() with timeout=0 should return immediately."""
+    consumer = TestConsumer(
+        {
+            'group.id': 'test-accumulate-zero',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 1000,
+            'auto.offset.reset': 'latest',
+        }
+    )
+    consumer.subscribe(['test-accumulate-topic'])
+
+    start = time.time()
+    msglist = consumer.consume(num_messages=10, timeout=0.0)
+    elapsed = time.time() - start
+
+    assert isinstance(msglist, list), "consume() should return a list"
+    assert len(msglist) == 0, "Expected empty list"
+    assert elapsed < WAKEABLE_POLL_TIMEOUT_MIN, f"Zero timeout should return immediately, took {elapsed:.2f}s"
+    consumer.close()

From 3f7925b2a7aa9d6f8d7692f5b2f5cdead230df49 Mon Sep 17 00:00:00 2001
From: Kaushik Raina <kraina@confluent.io>
Date: Mon, 20 Apr 2026 11:03:43 +0530
Subject: [PATCH 2/6] Minor fixes

---
 src/confluent_kafka/src/Consumer.c            | 20 +++++++++----------
 .../test_consumer_wakeable_poll_consume.py    |  4 ++--
 tests/test_Wakeable.py                        |  6 ++++--
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/src/confluent_kafka/src/Consumer.c b/src/confluent_kafka/src/Consumer.c
index f83fbcf9a..a1eb53cf9 100644
--- a/src/confluent_kafka/src/Consumer.c
+++ b/src/confluent_kafka/src/Consumer.c
@@ -1161,12 +1161,11 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
                     rkqu, total_timeout_ms, rkmessages, num_messages);
 
                 if (msgs_received_count < 0) {
-                        /* Error - need to restore GIL before setting error */
-                        PyEval_RestoreThread(cs.thread_state);
+                        if (CallState_end(self, &cs))
+                                cfl_PyErr_Format(
+                                    rd_kafka_last_error(), "%s",
+                                    rd_kafka_err2str(rd_kafka_last_error()));
                         free(rkmessages);
-                        cfl_PyErr_Format(
-                            rd_kafka_last_error(), "%s",
-                            rd_kafka_err2str(rd_kafka_last_error()));
                         return NULL;
                 }
         } else {
@@ -1189,15 +1188,14 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
                                     (unsigned int)msgs_received_count);
 
                         if (chunk_msg_count < 0) {
-                                /* Error - destroy accumulated messages,
-                                 * restore GIL, and raise */
                                 for (i = 0; i < msgs_received_count; i++)
                                         rd_kafka_message_destroy(rkmessages[i]);
-                                PyEval_RestoreThread(cs.thread_state);
+                                if (CallState_end(self, &cs))
+                                        cfl_PyErr_Format(
+                                            rd_kafka_last_error(), "%s",
+                                            rd_kafka_err2str(
+                                                rd_kafka_last_error()));
                                 free(rkmessages);
-                                cfl_PyErr_Format(
-                                    rd_kafka_last_error(), "%s",
-                                    rd_kafka_err2str(rd_kafka_last_error()));
                                 return NULL;
                         }
 
diff --git a/tests/integration/consumer/test_consumer_wakeable_poll_consume.py b/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
index ab49627b1..154d72fbd 100644
--- a/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
+++ b/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
@@ -270,8 +270,8 @@ def produce_in_waves():
     msglist = consumer.consume(num_messages=10, timeout=10.0)
 
     # Should have accumulated messages across multiple waves
-    assert len(msglist) >= 10, (
-        f"Expected at least 10 messages accumulated across waves, got {len(msglist)}. "
+    assert len(msglist) == 10, (
+        f"Expected exactly 10 messages accumulated across waves, got {len(msglist)}. "
         f"consume() may be returning early on the first wave."
     )
 
diff --git a/tests/test_Wakeable.py b/tests/test_Wakeable.py
index 3f962f38d..9f57988a5 100644
--- a/tests/test_Wakeable.py
+++ b/tests/test_Wakeable.py
@@ -1540,7 +1540,9 @@ def test_consume_accumulation_short_timeout_skips_chunking():
 
     assert isinstance(msglist, list), "consume() should return a list"
     assert len(msglist) == 0, "Expected empty list"
-    assert elapsed < WAKEABLE_POLL_TIMEOUT_MIN, f"Short timeout should not use chunking, took {elapsed:.2f}s"
+    assert (
+        elapsed <= WAKEABLE_POLL_TIMEOUT_MAX
+    ), f"Short timeout should not behave like a long chunked wait, took {elapsed:.2f}s"
     consumer.close()
 
 
@@ -1562,5 +1564,5 @@ def test_consume_accumulation_zero_timeout_nonblocking():
 
     assert isinstance(msglist, list), "consume() should return a list"
     assert len(msglist) == 0, "Expected empty list"
-    assert elapsed < WAKEABLE_POLL_TIMEOUT_MIN, f"Zero timeout should return immediately, took {elapsed:.2f}s"
+    assert elapsed <= WAKEABLE_POLL_TIMEOUT_MAX, f"Zero timeout should return immediately, took {elapsed:.2f}s"
     consumer.close()

From f218dc44933299524eb53e8dc9c82de41bd7347a Mon Sep 17 00:00:00 2001
From: Kaushik Raina <kraina@confluent.io>
Date: Mon, 20 Apr 2026 14:50:03 +0530
Subject: [PATCH 3/6] Minor fix

---
 src/confluent_kafka/src/Producer.c            |  5 ++-
 .../test_producer_wakeable_poll_flush.py      | 35 +++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/confluent_kafka/src/Producer.c b/src/confluent_kafka/src/Producer.c
index 0bfc810b8..4acf8baee 100644
--- a/src/confluent_kafka/src/Producer.c
+++ b/src/confluent_kafka/src/Producer.c
@@ -394,7 +394,10 @@ static int Producer_poll0(Handle *self, int tmout) {
                                 r = chunk_result;
                                 break;
                         }
-                        r += chunk_result; /* Accumulate events processed */
+                        r += chunk_result;
+
+                        if (chunk_result > 0)
+                                break;
 
                         chunk_count++;
 
diff --git a/tests/integration/producer/test_producer_wakeable_poll_flush.py b/tests/integration/producer/test_producer_wakeable_poll_flush.py
index 49d1d47d0..b72f00eed 100644
--- a/tests/integration/producer/test_producer_wakeable_poll_flush.py
+++ b/tests/integration/producer/test_producer_wakeable_poll_flush.py
@@ -119,6 +119,41 @@ def delivery_callback(err, msg):
     consumer.close()
 
 
+def test_poll_returns_early_after_delivery_callback(kafka_cluster):
+    """Test that poll() returns early after delivery callback fires."""
+    topic = kafka_cluster.create_topic_and_wait_propogation('test-poll-early-return')
+
+    delivery_called = []
+
+    def delivery_callback(err, msg):
+        delivery_called.append(time.time())
+
+    producer_conf = kafka_cluster.client_conf(
+        {
+            'socket.timeout.ms': 100,
+            'message.timeout.ms': 10000,
+        }
+    )
+    producer = kafka_cluster.cimpl_producer(producer_conf)
+
+    producer.produce(topic, value=b'early-return-test', on_delivery=delivery_callback)
+
+    # Poll with a long timeout — should return early once callback fires
+    poll_timeout = 5.0
+    start = time.time()
+    events = producer.poll(timeout=poll_timeout)
+    elapsed = time.time() - start
+
+    assert len(delivery_called) == 1, "Expected delivery callback to fire"
+    assert events > 0, "Expected at least 1 event served"
+    assert elapsed < poll_timeout - 1.0, (
+        f"poll({poll_timeout}) took {elapsed:.2f}s — should have returned "
+        f"early after delivery callback, not blocked for full timeout"
+    )
+
+    producer.close()
+
+
 def test_flush_message_delivery_with_wakeable_pattern(kafka_cluster):
     """Test that flush() correctly delivers messages when using wakeable pattern.
 

From bd11dd5f4654aba63b4485d5a894a8637cb9d60b Mon Sep 17 00:00:00 2001
From: Kaushik Raina <kraina@confluent.io>
Date: Tue, 14 Jul 2026 14:40:40 +0530
Subject: [PATCH 4/6] Remove wakeability from consumer.consume function

---
 src/confluent_kafka/src/Consumer.c            | 118 +++---------
 src/confluent_kafka/src/Producer.c            |   5 +-
 .../test_consumer_wakeable_poll_consume.py    |  19 +-
 tests/test_Wakeable.py                        | 181 ++++--------------
 4 files changed, 82 insertions(+), 241 deletions(-)

diff --git a/src/confluent_kafka/src/Consumer.c b/src/confluent_kafka/src/Consumer.c
index a1eb53cf9..67f00d4ea 100644
--- a/src/confluent_kafka/src/Consumer.c
+++ b/src/confluent_kafka/src/Consumer.c
@@ -1090,15 +1090,14 @@ Consumer_memberid(Handle *self, PyObject *args, PyObject *kwargs) {
 }
 
 /**
- * @brief Consume a batch of messages from the subscribed topics.
+ * @brief Consume a batch of up to num_messages messages from the subscribed
+ * topics.
  *
- * Instead of a single blocking call to rd_kafka_consume_batch_queue() with the
- * full timeout, this function:
- * 1. Splits the timeout into 200ms chunks
- * 2. Calls rd_kafka_consume_batch_queue() with chunk timeout
- * 3. Between chunks, re-acquires GIL and calls PyErr_CheckSignals()
- * 4. If signal detected, returns NULL (raises KeyboardInterrupt)
- * 5. Continues until messages received, timeout expired, or signal detected.
+ * This makes a single blocking call to rd_kafka_consume_batch_queue(), which
+ * gathers up to num_messages messages, blocking for up to the full timeout.
+ * The call is not interruptible while it is blocked: a pending signal such as
+ * Ctrl+C is only observed once the call returns, so with the default infinite
+ * timeout an idle consume() can block until a message arrives.
  *
  * @param self Consumer handle
  * @param args Positional arguments (unused)
@@ -1107,8 +1106,9 @@ Consumer_memberid(Handle *self, PyObject *args, PyObject *kwargs) {
  *                consume per call. Default: 1. Maximum: 1000000.
  *              - timeout (float, optional): Timeout in seconds.
  *                Default: -1.0 (infinite timeout)
- * @return PyObject* List of Message objects, empty list if timeout, or NULL on
- * error (raises KeyboardInterrupt if signal detected)
+ * @return PyObject* List of up to num_messages Message objects, empty list if
+ * the timeout elapses with none available, or NULL on error or when a pending
+ * signal (e.g. KeyboardInterrupt) is raised after the call returns.
  */
 static PyObject *
 Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
@@ -1119,12 +1119,7 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
         PyObject *msglist;
         rd_kafka_queue_t *rkqu = self->u.Consumer.rkqu;
         CallState cs;
-        Py_ssize_t i, msgs_received_count = 0;
-        Py_ssize_t chunk_msg_count;
-        const int CHUNK_TIMEOUT_MS = 200; /* 200ms chunks for signal checking */
-        int total_timeout_ms;
-        int chunk_timeout_ms;
-        int chunk_count = 0;
+        Py_ssize_t i, n;
 
         if (!self->rk) {
                 PyErr_SetString(PyExc_RuntimeError, ERR_MSG_CONSUMER_CLOSED);
@@ -1142,8 +1137,6 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
                 return NULL;
         }
 
-        total_timeout_ms = cfl_timeout_ms(tmout);
-
         rkmessages = malloc(num_messages * sizeof(rd_kafka_message_t *));
         if (!rkmessages) {
                 PyErr_NoMemory();
@@ -1152,85 +1145,27 @@ Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
 
         CallState_begin(self, &cs);
 
-        /* Skip wakeable poll pattern for non-blocking or very short timeouts.
-         * This avoids unnecessary GIL re-acquisition that can interfere with
-         * ThreadPool. Only use wakeable poll for
-         * blocking calls that need to be interruptible. */
-        if (total_timeout_ms >= 0 && total_timeout_ms < CHUNK_TIMEOUT_MS) {
-                msgs_received_count = (Py_ssize_t)rd_kafka_consume_batch_queue(
-                    rkqu, total_timeout_ms, rkmessages, num_messages);
-
-                if (msgs_received_count < 0) {
-                        if (CallState_end(self, &cs))
-                                cfl_PyErr_Format(
-                                    rd_kafka_last_error(), "%s",
-                                    rd_kafka_err2str(rd_kafka_last_error()));
-                        free(rkmessages);
-                        return NULL;
-                }
-        } else {
-                while (1) {
-                        /* Calculate timeout for this chunk */
-                        chunk_timeout_ms = calculate_chunk_timeout(
-                            total_timeout_ms, chunk_count, CHUNK_TIMEOUT_MS);
-                        if (chunk_timeout_ms == 0) {
-                                /* Timeout expired */
-                                break;
-                        }
-
-                        /* Consume with chunk timeout, appending after
-                         * already-accumulated messages */
-                        chunk_msg_count =
-                            (Py_ssize_t)rd_kafka_consume_batch_queue(
-                                rkqu, chunk_timeout_ms,
-                                rkmessages + msgs_received_count,
-                                num_messages -
-                                    (unsigned int)msgs_received_count);
-
-                        if (chunk_msg_count < 0) {
-                                for (i = 0; i < msgs_received_count; i++)
-                                        rd_kafka_message_destroy(rkmessages[i]);
-                                if (CallState_end(self, &cs))
-                                        cfl_PyErr_Format(
-                                            rd_kafka_last_error(), "%s",
-                                            rd_kafka_err2str(
-                                                rd_kafka_last_error()));
-                                free(rkmessages);
-                                return NULL;
-                        }
-
-                        msgs_received_count += chunk_msg_count;
-
-                        /* If we got all requested messages, exit the loop */
-                        if (msgs_received_count >= (Py_ssize_t)num_messages) {
-                                break;
-                        }
-
-                        chunk_count++;
-
-                        /* Check for signals between chunks */
-                        if (check_signals_between_chunks(self, &cs)) {
-                                for (i = 0; i < msgs_received_count; i++)
-                                        rd_kafka_message_destroy(rkmessages[i]);
-                                free(rkmessages);
-                                return NULL;
-                        }
-                }
-        }
+        n = (Py_ssize_t)rd_kafka_consume_batch_queue(
+            rkqu, cfl_timeout_ms(tmout), rkmessages, num_messages);
 
-        /* Final GIL restore and signal check */
         if (!CallState_end(self, &cs)) {
-                for (i = 0; i < msgs_received_count; i++) {
+                for (i = 0; i < n; i++) {
                         rd_kafka_message_destroy(rkmessages[i]);
                 }
                 free(rkmessages);
                 return NULL;
         }
 
-        /* Create Python list from messages  */
-        msglist = PyList_New(msgs_received_count);
+        if (n < 0) {
+                free(rkmessages);
+                cfl_PyErr_Format(rd_kafka_last_error(), "%s",
+                                 rd_kafka_err2str(rd_kafka_last_error()));
+                return NULL;
+        }
+
+        msglist = PyList_New(n);
 
-        for (i = 0; i < msgs_received_count; i++) {
+        for (i = 0; i < n; i++) {
                 PyObject *msgobj = Message_new0(self, rkmessages[i]);
 #ifdef RD_KAFKA_V_HEADERS
                 /** Have to detach headers outside Message_new0 because it
@@ -1403,6 +1338,13 @@ static PyMethodDef Consumer_methods[] = {
      "  .. note: Callbacks may be called from this method, "
      "such as ``on_assign``, ``on_revoke``, et.al.\n"
      "\n"
+     "  .. note:: This is a blocking call and does not respond to signals "
+     "(e.g. Ctrl+C / SIGINT) while it is waiting; a pending signal is only "
+     "delivered once the call returns. With the default infinite timeout an "
+     "idle ``consume()`` can block indefinitely. Applications that need to "
+     "stay responsive to interrupts should pass a short ``timeout`` and call "
+     "``consume()`` in a loop, handling signals between calls.\n"
+     "\n"
      "  :param int num_messages: The maximum number of messages to return "
      "(default: 1).\n"
      "  :param float timeout: The maximum time to block waiting for message, "
diff --git a/src/confluent_kafka/src/Producer.c b/src/confluent_kafka/src/Producer.c
index 4acf8baee..726262fb4 100644
--- a/src/confluent_kafka/src/Producer.c
+++ b/src/confluent_kafka/src/Producer.c
@@ -353,7 +353,8 @@ Producer_produce(Handle *self, PyObject *args, PyObject *kwargs) {
  * 2. Calls rd_kafka_poll() with chunk timeout
  * 3. Between chunks, re-acquires GIL and calls PyErr_CheckSignals()
  * 4. If signal detected, returns -1 (raises KeyboardInterrupt)
- * 5. Continues until events processed, timeout expired, or signal detected
+ * 5. Returns as soon as a chunk serves any events, the timeout expires, or a
+ *    signal is detected
  *
  * @param self Producer handle
  * @param tmout Timeout in milliseconds (-1 for infinite)
@@ -394,7 +395,7 @@ static int Producer_poll0(Handle *self, int tmout) {
                                 r = chunk_result;
                                 break;
                         }
-                        r += chunk_result;
+                        r = chunk_result;
 
                         if (chunk_result > 0)
                                 break;
diff --git a/tests/integration/consumer/test_consumer_wakeable_poll_consume.py b/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
index 154d72fbd..4cddd962d 100644
--- a/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
+++ b/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
@@ -144,10 +144,9 @@ def test_consume_message_delivery_with_wakeable_pattern(kafka_cluster):
     consumer.close()
 
 
-def test_consume_accumulates_messages_across_chunks(kafka_cluster):
-    """Test that consume() accumulates messages across 200ms chunks.
-    This verifies that consume() doesn't return early on the first chunk of messages
-    when using the wakeable pattern.
+def test_consume_accumulates_messages_up_to_num_messages(kafka_cluster):
+    """Test that consume() gathers up to num_messages within the timeout,
+    rather than returning as soon as the first messages become available.
     """
     topic = kafka_cluster.create_topic_and_wait_propogation('test-consume-accumulate-chunks')
 
@@ -173,13 +172,12 @@ def test_consume_accumulates_messages_across_chunks(kafka_cluster):
     # Wait for subscription and partition assignment
     time.sleep(2.0)
 
-    # Consume with num_messages=10 and a generous timeout.
-    # Before the fix: would return < 10 (whatever arrived in the first 200ms chunk)
-    # After the fix: accumulates across chunks until 10 are collected
+    # Consume num_messages=10 with a generous timeout; the batch call gathers
+    # all 10 within the timeout rather than returning on the first few.
     msglist = consumer.consume(num_messages=num_produced, timeout=10.0)
 
     assert len(msglist) == num_produced, (
-        f"Expected {num_produced} messages but got {len(msglist)}. " f"consume() may not be accumulating across chunks."
+        f"Expected {num_produced} messages but got {len(msglist)}. " f"consume() did not gather the full batch within the timeout."
     )
 
     for i, msg in enumerate(msglist):
@@ -230,9 +228,8 @@ def test_consume_returns_partial_on_timeout(kafka_cluster):
 
 
 def test_consume_accumulates_messages_produced_in_waves(kafka_cluster):
-    """Test that consume() accumulates messages that arrive in multiple waves.
-    This verifies that consume() doesn't return early on the first wave of messages
-    when using the wakeable pattern.
+    """Test that consume() accumulates messages that arrive in multiple waves,
+    rather than returning as soon as the first wave becomes available.
     """
     import threading
 
diff --git a/tests/test_Wakeable.py b/tests/test_Wakeable.py
index 9f57988a5..4d62c0160 100644
--- a/tests/test_Wakeable.py
+++ b/tests/test_Wakeable.py
@@ -47,7 +47,10 @@
 #
 # Consumer Implementation (Consumer.c):
 #   - Consumer.poll() uses wakeable pattern for timeouts >= 200ms
-#   - Consumer.consume() uses wakeable pattern for timeouts >= 200ms
+#   - Consumer.consume() is intentionally NOT wakeable: it makes a single
+#     blocking batch call so offset storage stays atomic with the batch the
+#     caller actually receives. It is therefore not interruptible mid-call and
+#     has no wakeability/interruptibility tests here.
 #
 # How We Test Wakeability:
 # ------------------------
@@ -783,88 +786,17 @@ def test_consumer_wakeable_poll_edge_cases():
     consumer4.close()
 
 
-def test_consumer_wakeable_consume_interruptibility_and_messages():
-    """Test consume() interruptibility (main fix) and message handling."""
-    topic = 'test-consume-interrupt-topic'
+def test_consumer_consume_timeout_and_message_handling():
+    """Test consume() batch timeout behavior and message handling.
 
-    # Assertion 1: Infinite timeout can be interrupted immediately
-    consumer1 = TestConsumer(
-        {
-            'group.id': 'test-consume-infinite-immediate',
-            'socket.timeout.ms': 100,
-            'session.timeout.ms': 1000,
-            'auto.offset.reset': 'latest',
-        }
-    )
-    consumer1.subscribe([topic])
-
-    interrupt_thread = threading.Thread(target=lambda: TestUtils.send_sigint_after_delay(0.1))
-    interrupt_thread.daemon = True
-    interrupt_thread.start()
-
-    interrupted = False
-    try:
-        consumer1.consume()  # Infinite timeout, default num_messages=1
-    except KeyboardInterrupt:
-        interrupted = True
-    finally:
-        consumer1.close()
-
-    assert interrupted, "Assertion 1 failed: Should have raised KeyboardInterrupt"
-
-    # Assertion 2: Finite timeout can be interrupted before timeout expires
-    consumer2 = TestConsumer(
-        {
-            'group.id': 'test-consume-finite-interrupt',
-            'socket.timeout.ms': 100,
-            'session.timeout.ms': 1000,
-            'auto.offset.reset': 'latest',
-        }
-    )
-    consumer2.subscribe([topic])
-
-    interrupt_thread = threading.Thread(target=lambda: TestUtils.send_sigint_after_delay(0.3))
-    interrupt_thread.daemon = True
-    interrupt_thread.start()
-
-    interrupted = False
-    timeout_value = WAKEABLE_POLL_TIMEOUT_MAX  # Use constant instead of hardcoded 2.0
-    try:
-        consumer2.consume(num_messages=10, timeout=timeout_value)  # Use constant for timeout
-    except KeyboardInterrupt:
-        interrupted = True
-    finally:
-        consumer2.close()
-
-    assert interrupted, "Assertion 2 failed: Should have raised KeyboardInterrupt"
-
-    # Assertion 3: Signal sent after multiple chunks still interrupts quickly
-    consumer3 = TestConsumer(
-        {
-            'group.id': 'test-consume-multiple-chunks',
-            'socket.timeout.ms': 100,
-            'session.timeout.ms': 1000,
-            'auto.offset.reset': 'latest',
-        }
-    )
-    consumer3.subscribe([topic])
-
-    interrupt_thread = threading.Thread(target=lambda: TestUtils.send_sigint_after_delay(0.6))
-    interrupt_thread.daemon = True
-    interrupt_thread.start()
-
-    interrupted = False
-    try:
-        consumer3.consume(num_messages=5)  # Infinite timeout
-    except KeyboardInterrupt:
-        interrupted = True
-    finally:
-        consumer3.close()
-
-    assert interrupted, "Assertion 3 failed: Should have raised KeyboardInterrupt"
+    consume() is not wakeable/interruptible (single blocking batch call), so
+    this only covers the non-signal behavior: honoring the timeout and the
+    num_messages=0 short-circuit.
+    """
+    topic = 'test-consume-topic'
 
-    # Assertion 4: No signal - timeout works normally, returns empty list
-    consumer4 = TestConsumer(
+    # Assertion 1: No signal - timeout works normally, returns empty list
+    consumer1 = TestConsumer(
         {
             'group.id': 'test-consume-timeout-normal',
             'socket.timeout.ms': 100,
@@ -872,21 +804,21 @@ def test_consumer_wakeable_consume_interruptibility_and_messages():
             'auto.offset.reset': 'latest',
         }
     )
-    consumer4.subscribe([topic])
+    consumer1.subscribe([topic])
 
     start = time.time()
-    msglist = consumer4.consume(num_messages=10, timeout=0.5)  # 500ms, no signal
+    msglist = consumer1.consume(num_messages=10, timeout=0.5)  # 500ms, no signal
     elapsed = time.time() - start
 
-    assert isinstance(msglist, list), "Assertion 4 failed: consume() should return a list"
-    assert len(msglist) == 0, f"Assertion 4 failed: Expected empty list (timeout), got {len(msglist)} messages"
+    assert isinstance(msglist, list), "Assertion 1 failed: consume() should return a list"
+    assert len(msglist) == 0, f"Assertion 1 failed: Expected empty list (timeout), got {len(msglist)} messages"
     assert (
         WAKEABLE_POLL_TIMEOUT_MIN <= elapsed <= WAKEABLE_POLL_TIMEOUT_MAX
-    ), f"Assertion 4 failed: Normal timeout took {elapsed:.2f}s, expected ~0.5s"
-    consumer4.close()
+    ), f"Assertion 1 failed: Normal timeout took {elapsed:.2f}s, expected ~0.5s"
+    consumer1.close()
 
-    # Assertion 5: num_messages=0 returns empty list immediately
-    consumer5 = TestConsumer(
+    # Assertion 2: num_messages=0 returns empty list immediately
+    consumer2 = TestConsumer(
         {
             'group.id': 'test-consume-zero-messages',
             'socket.timeout.ms': 100,
@@ -894,18 +826,18 @@ def test_consumer_wakeable_consume_interruptibility_and_messages():
             'auto.offset.reset': 'latest',
         }
     )
-    consumer5.subscribe([topic])
+    consumer2.subscribe([topic])
 
     start = time.time()
-    msglist = consumer5.consume(num_messages=0, timeout=1.0)
+    msglist = consumer2.consume(num_messages=0, timeout=1.0)
     elapsed = time.time() - start
 
-    assert isinstance(msglist, list), "Assertion 5 failed: consume() should return a list"
-    assert len(msglist) == 0, "Assertion 5 failed: num_messages=0 should return empty list"
+    assert isinstance(msglist, list), "Assertion 2 failed: consume() should return a list"
+    assert len(msglist) == 0, "Assertion 2 failed: num_messages=0 should return empty list"
     assert (
         elapsed < WAKEABLE_POLL_TIMEOUT_MAX
-    ), f"Assertion 5 failed: num_messages=0 took {elapsed:.2f}s, expected < {WAKEABLE_POLL_TIMEOUT_MAX}s"
-    consumer5.close()
+    ), f"Assertion 2 failed: num_messages=0 took {elapsed:.2f}s, expected < {WAKEABLE_POLL_TIMEOUT_MAX}s"
+    consumer2.close()
 
 
 def test_consumer_wakeable_consume_edge_cases():
@@ -1319,11 +1251,14 @@ def blocking_call(t):
         ("producer", "poll"),
         ("producer", "flush"),
         ("consumer", "poll"),
-        ("consumer", "consume"),
     ],
 )
 def test_can_be_interrupted(api_type, method):
-    """Test that blocking operations can be interrupted."""
+    """Test that blocking operations can be interrupted.
+
+    consumer.consume() is intentionally excluded: it is a single blocking batch
+    call and is not interruptible.
+    """
     if api_type == "producer":
         obj = Producer({'bootstrap.servers': 'localhost:9092', 'socket.timeout.ms': 100, 'message.timeout.ms': 10})
         if method == "poll":
@@ -1347,15 +1282,9 @@ def blocking_call():
             }
         )
         obj.subscribe(['test-topic'])
-        if method == "poll":
-
-            def blocking_call():
-                return obj.poll()
 
-        else:  # consume
-
-            def blocking_call():
-                return obj.consume()
+        def blocking_call():
+            return obj.poll()
 
     interrupt_thread = threading.Thread(target=lambda: TestUtils.send_sigint_after_delay(0.1))
     interrupt_thread.daemon = True
@@ -1440,13 +1369,13 @@ def test_flush_empty_queue_returns_immediately():
     assert elapsed < WAKEABLE_POLL_TIMEOUT_MAX, f"Empty flush should return quickly, took {elapsed:.2f}s"
 
 
-# These tests verify that consume() correctly accumulates messages across
-# 200ms chunks instead of returning on the first non-empty chunk.
+# These tests verify consume()'s single blocking batch call: it honors the
+# timeout and returns up to num_messages messages (empty here, no broker).
 
 
-def test_consume_accumulates_across_chunks_no_messages():
+def test_consume_no_messages_waits_full_timeout():
     """consume() with no messages available should wait the full timeout
-    and return an empty list, not return early on the first empty chunk."""
+    and return an empty list."""
     consumer = TestConsumer(
         {
             'group.id': 'test-accumulate-no-msgs',
@@ -1469,34 +1398,6 @@ def test_consume_accumulates_across_chunks_no_messages():
     consumer.close()
 
 
-def test_consume_accumulation_signal_interrupts_cleanly():
-    """When Ctrl+C arrives during accumulation, consume() should raise
-    KeyboardInterrupt without leaking messages."""
-    consumer = TestConsumer(
-        {
-            'group.id': 'test-accumulate-signal',
-            'socket.timeout.ms': 100,
-            'session.timeout.ms': 1000,
-            'auto.offset.reset': 'latest',
-        }
-    )
-    consumer.subscribe(['test-accumulate-topic'])
-
-    interrupt_thread = threading.Thread(target=lambda: TestUtils.send_sigint_after_delay(0.3))
-    interrupt_thread.daemon = True
-    interrupt_thread.start()
-
-    interrupted = False
-    try:
-        consumer.consume(num_messages=100, timeout=WAKEABLE_POLL_TIMEOUT_MAX)
-    except KeyboardInterrupt:
-        interrupted = True
-    finally:
-        consumer.close()
-
-    assert interrupted, "Should have raised KeyboardInterrupt"
-
-
 def test_consume_accumulation_timeout_returns_partial():
     """consume() should return whatever messages it has when timeout expires,
     even if fewer than num_messages. With no broker, this means empty list."""
@@ -1522,8 +1423,8 @@ def test_consume_accumulation_timeout_returns_partial():
     consumer.close()
 
 
-def test_consume_accumulation_short_timeout_skips_chunking():
-    """consume() with timeout < 200ms should skip the chunked loop entirely."""
+def test_consume_short_timeout_returns_quickly():
+    """consume() with a short timeout should return promptly with an empty list."""
     consumer = TestConsumer(
         {
             'group.id': 'test-accumulate-short',
@@ -1542,7 +1443,7 @@ def test_consume_accumulation_short_timeout_skips_chunking():
     assert len(msglist) == 0, "Expected empty list"
     assert (
         elapsed <= WAKEABLE_POLL_TIMEOUT_MAX
-    ), f"Short timeout should not behave like a long chunked wait, took {elapsed:.2f}s"
+    ), f"Short timeout should return quickly, took {elapsed:.2f}s"
     consumer.close()
 
 

From 519ad64d75d42dcb71f7075610b12b77e299848d Mon Sep 17 00:00:00 2001
From: Kaushik Raina <kraina@confluent.io>
Date: Tue, 14 Jul 2026 16:39:27 +0530
Subject: [PATCH 5/6] Style fix

---
 .../consumer/test_consumer_wakeable_poll_consume.py           | 3 ++-
 tests/test_Wakeable.py                                        | 4 +---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/tests/integration/consumer/test_consumer_wakeable_poll_consume.py b/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
index 4cddd962d..2ad3a359e 100644
--- a/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
+++ b/tests/integration/consumer/test_consumer_wakeable_poll_consume.py
@@ -177,7 +177,8 @@ def test_consume_accumulates_messages_up_to_num_messages(kafka_cluster):
     msglist = consumer.consume(num_messages=num_produced, timeout=10.0)
 
     assert len(msglist) == num_produced, (
-        f"Expected {num_produced} messages but got {len(msglist)}. " f"consume() did not gather the full batch within the timeout."
+        f"Expected {num_produced} messages but got {len(msglist)}. "
+        f"consume() did not gather the full batch within the timeout."
     )
 
     for i, msg in enumerate(msglist):
diff --git a/tests/test_Wakeable.py b/tests/test_Wakeable.py
index 4d62c0160..0d53f654a 100644
--- a/tests/test_Wakeable.py
+++ b/tests/test_Wakeable.py
@@ -1441,9 +1441,7 @@ def test_consume_short_timeout_returns_quickly():
 
     assert isinstance(msglist, list), "consume() should return a list"
     assert len(msglist) == 0, "Expected empty list"
-    assert (
-        elapsed <= WAKEABLE_POLL_TIMEOUT_MAX
-    ), f"Short timeout should return quickly, took {elapsed:.2f}s"
+    assert elapsed <= WAKEABLE_POLL_TIMEOUT_MAX, f"Short timeout should return quickly, took {elapsed:.2f}s"
     consumer.close()
 
 

From 6c679bec0c82c43b54567cceabebaaba2f5f3c1a Mon Sep 17 00:00:00 2001
From: Kaushik Raina <kraina@confluent.io>
Date: Thu, 16 Jul 2026 15:27:24 +0530
Subject: [PATCH 6/6] Fix minor

---
 tests/test_Wakeable.py | 20 ++++++++------------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/tests/test_Wakeable.py b/tests/test_Wakeable.py
index 0d53f654a..088da078f 100644
--- a/tests/test_Wakeable.py
+++ b/tests/test_Wakeable.py
@@ -1249,7 +1249,6 @@ def blocking_call(t):
     "api_type,method",
     [
         ("producer", "poll"),
-        ("producer", "flush"),
         ("consumer", "poll"),
     ],
 )
@@ -1261,16 +1260,9 @@ def test_can_be_interrupted(api_type, method):
     """
     if api_type == "producer":
         obj = Producer({'bootstrap.servers': 'localhost:9092', 'socket.timeout.ms': 100, 'message.timeout.ms': 10})
-        if method == "poll":
-
-            def blocking_call():
-                return obj.poll()
 
-        else:  # flush
-            obj.produce('test-topic', value='test', callback=lambda err, msg: None)
-
-            def blocking_call():
-                return obj.flush()
+        def blocking_call():
+            return obj.poll()
 
     else:  # consumer
         obj = TestConsumer(
@@ -1296,8 +1288,12 @@ def blocking_call():
     except KeyboardInterrupt:
         interrupted = True
     finally:
-        # Wait for signal thread to complete
-        time.sleep(0.2)
+        # A SIGINT delivered after the blocking call already returned must not
+        # escape cleanup and abort the whole pytest session.
+        try:
+            time.sleep(0.2)
+        except KeyboardInterrupt:
+            interrupted = True
         obj.close()
 
     # Key assertion: operation was interruptible