From 23d43feb19238d2d6c8985f30f1d6f6412271a18 Mon Sep 17 00:00:00 2001
From: Shunping Huang <shunping@google.com>
Date: Tue, 30 Jun 2026 22:11:44 -0400
Subject: [PATCH] Acknowledge messages asap in pubsub test

---
 .../io/gcp/pubsub_integration_test.py         | 20 ++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
index 12b200764578..89fd4461beb3 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
@@ -363,7 +363,6 @@ def test_batch_write_with_ordering_key(self):
       # Retry pulling to handle PubSub delivery delays
       received_messages = []
       received_message_ids = set()
-      ack_ids = []
       deadline = time.time() + 60  # wait up to 60 seconds
       while time.time() < deadline:
         response = self.sub_client.pull(
@@ -371,8 +370,18 @@ def test_batch_write_with_ordering_key(self):
                 'subscription': ordering_sub.name,
                 'max_messages': 10,
             })
+        # Acknowledge received messages immediately. This is critical when message
+        # ordering is enabled because outstanding (unacknowledged) messages block
+        # the delivery of subsequent messages with the same ordering key.
+        ack_ids = [msg.ack_id for msg in response.received_messages]
+        if ack_ids:
+          self.sub_client.acknowledge(
+              request={
+                  'subscription': ordering_sub.name,
+                  'ack_ids': ack_ids,
+              })
+
         for msg in response.received_messages:
-          ack_ids.append(msg.ack_id)
           # Pub/Sub guarantees at-least-once delivery, so we must deduplicate
           # messages by message_id to handle potential duplicate deliveries.
           if msg.message.message_id not in received_message_ids:
@@ -391,13 +400,6 @@ def test_batch_write_with_ordering_key(self):
       self.assertEqual(received_map[b'order_data001'].ordering_key, 'key1')
       self.assertEqual(received_map[b'order_data002'].ordering_key, 'key1')
       self.assertEqual(received_map[b'order_data003'].ordering_key, 'key2')
-
-      if ack_ids:
-        self.sub_client.acknowledge(
-            request={
-                'subscription': ordering_sub.name,
-                'ack_ids': ack_ids,
-            })
     finally:
       self.sub_client.delete_subscription(
           request={'subscription': ordering_sub.name})