Fix postcommit python flaky tests (#34536)

* Fix postcommit python flaky tests

* Fix formatting

* Fix formatting

* Fix PostCommit Python tests related comments

* Formatting

Author: akashorabek

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

@@ -131,6 +131,11 @@ def _setup_new_types_env(self):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=self.project, datasetId=self.dataset_id, table=table)
     self.bigquery_client.client.tables.Insert(request)
+
+    # Call get_table so that we wait until the table is visible.
+    _ = self.bigquery_client.get_table(
+        self.project, self.dataset_id, NEW_TYPES_INPUT_TABLE)
+
     table_data = [{
         'bytes': b'xyw', 'date': '2011-01-01', 'time': '23:59:59.999999'
     }, {

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

@@ -158,6 +158,8 @@ def create_table(cls, table_name):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=cls.project, datasetId=cls.dataset_id, table=table)
     cls.bigquery_client.client.tables.Insert(request)
+    # Call get_table so that we wait until the table is visible.
+    _ = cls.bigquery_client.get_table(cls.project, cls.dataset_id, table_name)
     cls.bigquery_client.insert_rows(
         cls.project, cls.dataset_id, table_name, cls.TABLE_DATA)
 
@@ -395,6 +397,8 @@ def _create_table(cls, table_name):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=cls.project, datasetId=cls.dataset_id, table=table)
     cls.bigquery_client.client.tables.Insert(request)
+    # Call get_table so that we wait until the table is visible.
+    _ = cls.bigquery_client.get_table(cls.project, cls.dataset_id, table_name)
     cls.bigquery_client.insert_rows(
         cls.project, cls.dataset_id, table_name, cls.TABLE_DATA)
 
@@ -656,6 +660,8 @@ def create_table(cls, table_name):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=cls.project, datasetId=cls.dataset_id, table=table)
     cls.bigquery_client.client.tables.Insert(request)
+    # Call get_table so that we wait until the table is visible.
+    _ = cls.bigquery_client.get_table(cls.project, cls.dataset_id, table_name)
     row_data = {
         'float': 0.33,
         'numeric': Decimal('10'),
@@ -760,8 +766,8 @@ def setUpClass(cls):
     cls.table_name2 = 'python_rd_table_2'
     cls.table_schema2 = cls.create_table(
         cls.table_name2, cls.TABLE_DATA_2, cls.SCHEMA_BQ)
-    table_id2 = '{}.{}'.format(cls.dataset_id, cls.table_name2)
-    cls.query2 = 'SELECT number, str FROM %s' % table_id2
+    cls.query2 = 'SELECT number, str FROM [%s:%s.%s]' % (
+        cls.project, cls.dataset_id, cls.table_name2)
 
     cls.table_name3 = 'python_rd_table_3'
     cls.table_schema3 = cls.create_table(
@@ -779,6 +785,8 @@ def create_table(cls, table_name, data, table_schema):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=cls.project, datasetId=cls.dataset_id, table=table)
     cls.bigquery_client.client.tables.Insert(request)
+    # Call get_table so that we wait until the table is visible.
+    _ = cls.bigquery_client.get_table(cls.project, cls.dataset_id, table_name)
     cls.bigquery_client.insert_rows(
         cls.project, cls.dataset_id, table_name, data)
     return table_schema

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

@@ -453,8 +453,12 @@ def test_big_query_write_insert_errors_reporting(self):
 
       assert_that(
           errors[BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS]
-          | 'ParseErrors' >> beam.Map(lambda err: (err[1], err[2])),
-          equal_to(bq_result_errors))
+          | 'ParseErrors' >> beam.Map(lambda err: (err[1], err[2]))
+          | 'ToList' >> beam.combiners.ToList()
+          | 'SortErrors' >> beam.Map(
+              lambda errs: sorted(errs, key=lambda x: x[0].get("number", 0))),
+          equal_to(
+              [sorted(bq_result_errors, key=lambda x: x[0].get("number", 0))]))
 
   @pytest.mark.it_postcommit
   def test_big_query_write_insert_non_transient_api_call_error(self):

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

@@ -29,6 +29,7 @@
 # pytype: skip-file
 
 import logging
+import time
 import unittest
 import uuid
 import zlib
@@ -58,6 +59,9 @@
 except ImportError:
   NotFound = None
 
+# Number of seconds to wait for bucket deletion to propagate.
+WAIT_DELETE_BUCKET_PROPAGATION_SECONDS = 10
+
 
 @unittest.skipIf(gcsio is None, 'GCP dependencies are not installed')
 @parameterized_class(
@@ -221,6 +225,7 @@ def test_create_default_bucket(self, mock_default_gcs_bucket_name):
     if existing_bucket:
       try:
         existing_bucket.delete()
+        time.sleep(WAIT_DELETE_BUCKET_PROPAGATION_SECONDS)
       except NotFound:
         # Bucket existence check from get_bucket may be inaccurate due to gcs
         # cache or delay

sdks/python/apache_beam/ml/gcp/recommendations_ai.py

@@ -33,6 +33,7 @@
 from apache_beam.transforms import ParDo
 from apache_beam.transforms import PTransform
 from apache_beam.transforms.util import GroupIntoBatches
+from apache_beam.utils import retry
 from cachetools.func import ttl_cache
 
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
@@ -53,6 +54,7 @@
 ]
 
 FAILED_CATALOG_ITEMS = "failed_catalog_items"
+MAX_RETRIES = 5
 
 
 @ttl_cache(maxsize=128, ttl=3600)
@@ -154,13 +156,19 @@ def process(self, element):
     request = recommendationengine.CreateCatalogItemRequest(
         parent=self.parent, catalog_item=catalog_item)
 
-    try:
-      created_catalog_item = self._client.create_catalog_item(
+    @retry.with_exponential_backoff(
+        num_retries=MAX_RETRIES,
+        retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter
+    )
+    def create_item():
+      return self._client.create_catalog_item(
           request=request,
           retry=self.retry,
           timeout=self.timeout,
           metadata=self.metadata)
 
+    try:
+      created_catalog_item = create_item()
       self.counter.inc()
       yield recommendationengine.CatalogItem.to_dict(created_catalog_item)
     except Exception: