Acknowledge messages asap in pubsub test (#39183)

Author: shunping

sdks/python/apache_beam/io/gcp/pubsub_integration_test.py

@@ -363,16 +363,25 @@ 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(
             request={
                 '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})