Minor fixes

Author: k-raina

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;
                         }
 

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."
     )
 

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()