Fix flaky Python tests

sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py

@@ -35,6 +35,7 @@
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.runners.runner import PipelineState
 from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.utils import retry
 
 # Protect against environments where bigtable library is not available.
 # pylint: disable=wrong-import-order, wrong-import-position
@@ -106,13 +107,13 @@ class BigtableIOWriteTest(unittest.TestCase):
 
   """
   DEFAULT_TABLE_PREFIX = "python-test"
-  instance_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
-  cluster_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
-  table_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
   number = 500
   LOCATION_ID = "us-east1-b"
 
   def setUp(self):
+    self.instance_id = self.DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
+    self.cluster_id = self.DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
+    self.table_id = self.DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]
     try:
       from google.cloud.bigtable import enums
       self.STORAGE_TYPE = enums.StorageType.HDD
@@ -136,8 +137,17 @@ def setUp(self):
           self.cluster_id,
           self.LOCATION_ID,
           default_storage_type=self.STORAGE_TYPE)
-      operation = self.instance.create(clusters=[cluster])
-      operation.result(timeout=300)  # Wait up to 5 min.
+
+      @retry.with_exponential_backoff(
+          num_retries=3,
+          initial_delay_secs=5,
+          retry_filter=retry.
+          retry_on_server_errors_timeout_or_quota_issues_filter)
+      def _create_instance():
+        operation = self.instance.create(clusters=[cluster])
+        operation.result(timeout=300)  # Wait up to 5 min.
+
+      _create_instance()
 
     self.table = self.instance.table(self.table_id)
 

sdks/python/apache_beam/examples/streaming_wordcount_it_test.py

@@ -66,12 +66,19 @@ def setUp(self):
             self.project, OUTPUT_SUB + self.uuid),
         topic=self.output_topic.name,
         ack_deadline_seconds=60)
+    import time
+    time.sleep(10)
 
   def _inject_numbers(self, topic, num_messages):
     """Inject numbers as test data to PubSub."""
     logging.debug('Injecting %d numbers to topic %s', num_messages, topic.name)
+    futures = []
     for n in range(num_messages):
-      self.pub_client.publish(self.input_topic.name, str(n).encode('utf-8'))
+      futures.append(
+          self.pub_client.publish(
+              self.input_topic.name, str(n).encode('utf-8')))
+    for future in futures:
+      future.result()
 
   def tearDown(self):
     test_utils.cleanup_subscriptions(
@@ -88,7 +95,7 @@ def test_streaming_wordcount_it(self):
     # Set extra options to the pipeline for test purpose
     state_verifier = PipelineStateMatcher(PipelineState.RUNNING)
     pubsub_msg_verifier = PubSubMessageMatcher(
-        self.project, self.output_sub.name, expected_msg, timeout=400)
+        self.project, self.output_sub.name, expected_msg, timeout=600)
     extra_opts = {
         'input_subscription': self.input_sub.name,
         'output_topic': self.output_topic.name,

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

@@ -290,6 +290,11 @@ def process(self, direct_row):
           direct_row._get_mutation_pbs() if hasattr(
               direct_row, '_get_mutation_pbs') else direct_row._get_mutations())
       for mutation in mutations:
+        if not hasattr(mutation, '__contains__'):
+          if hasattr(mutation, '_to_pb'):
+            mutation = mutation._to_pb()
+          elif hasattr(mutation, 'to_pb'):
+            mutation = mutation.to_pb()
         if mutation.__contains__("set_cell"):
           mutation_dict = {
               "type": b'SetCell',

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

@@ -122,11 +122,18 @@ class TestBigtableDirectRowToBeamRow(unittest.TestCase):
   @staticmethod
   def _get_mutation_pbs(direct_row):
     # In google-cloud-bigtable >= 2.44.0, _get_mutations() returns Python
-    # dataclass objects instead of protobuf messages; use _get_mutation_pbs()
-    # to retrieve Mutation protobuf messages.
-    if hasattr(direct_row, '_get_mutation_pbs'):
-      return direct_row._get_mutation_pbs()
-    return direct_row._get_mutations()
+    # dataclass objects instead of protobuf messages; convert each item via
+    # _to_pb() when needed.
+    mutations = (
+        direct_row._get_mutation_pbs() if hasattr(
+            direct_row, '_get_mutation_pbs') else direct_row._get_mutations())
+    return [
+        m._to_pb()
+        if not hasattr(m, '__contains__') and hasattr(m, '_to_pb') else (
+            m.to_pb()
+            if not hasattr(m, '__contains__') and hasattr(m, 'to_pb') else m)
+        for m in mutations
+    ]
 
   def test_set_cell(self):
     # create some set cell mutations

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

@@ -36,6 +36,7 @@
 from apache_beam.testing import test_utils
 from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher
 from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.utils import retry
 
 INPUT_TOPIC = 'psit_topic_input'
 OUTPUT_TOPIC = 'psit_topic_output'
@@ -138,20 +139,27 @@ def setUp(self):
     self.project = self.test_pipeline.get_option('project')
     self.uuid = str(uuid.uuid4())
 
-    # Set up PubSub environment.
+    # Set up PubSub environment with retries for transient 504/403 API errors.
     from google.cloud import pubsub
     self.pub_client = pubsub.PublisherClient()
-    self.input_topic = self.pub_client.create_topic(
+    self.sub_client = pubsub.SubscriberClient()
+
+    gcp_retry = retry.with_exponential_backoff(
+        num_retries=3,
+        initial_delay_secs=2,
+        retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter
+    )
+
+    self.input_topic = gcp_retry(self.pub_client.create_topic)(
         name=self.pub_client.topic_path(self.project, INPUT_TOPIC + self.uuid))
-    self.output_topic = self.pub_client.create_topic(
+    self.output_topic = gcp_retry(self.pub_client.create_topic)(
         name=self.pub_client.topic_path(self.project, OUTPUT_TOPIC + self.uuid))
 
-    self.sub_client = pubsub.SubscriberClient()
-    self.input_sub = self.sub_client.create_subscription(
+    self.input_sub = gcp_retry(self.sub_client.create_subscription)(
         name=self.sub_client.subscription_path(
             self.project, INPUT_SUB + self.uuid),
         topic=self.input_topic.name)
-    self.output_sub = self.sub_client.create_subscription(
+    self.output_sub = gcp_retry(self.sub_client.create_subscription)(
         name=self.sub_client.subscription_path(
             self.project, OUTPUT_SUB + self.uuid),
         topic=self.output_topic.name)

sdks/python/apache_beam/ml/anomaly/transforms_test.py

@@ -84,11 +84,14 @@ def _prediction_is_equal_to(a: AnomalyPrediction, b: AnomalyPrediction):
     return False
 
   if a.threshold != b.threshold:
-    return False
+    if not (a.threshold is not None and b.threshold is not None and
+            math.isclose(a.threshold, b.threshold, rel_tol=1e-6, abs_tol=1e-6)):
+      return False
 
   if a.score != b.score:
     if not (a.score is not None and b.score is not None and
-            math.isnan(a.score) and math.isnan(b.score)):
+            ((math.isnan(a.score) and math.isnan(b.score)) or
+             math.isclose(a.score, b.score, rel_tol=1e-6, abs_tol=1e-6))):
       return False
 
   if a.label != b.label:
@@ -246,8 +249,11 @@ class TestData:
   ])
   def test_one_detector(self, input, expected):
     detector = ZScore(features=["x1"], model_id="zscore_x1")
-    with TestPipeline() as p:
-      result = (p | beam.Create(input) | AnomalyDetection(detector))
+    with TestPipeline(
+        additional_pipeline_args=["--experiments=prism_disable_sdf_split"
+                                  ]) as p:
+      result = (
+          p | beam.Create(input, reshuffle=False) | AnomalyDetection(detector))
 
       if isinstance(input[0], tuple):
         assert_that(result, equal_to(expected, _keyed_result_is_equal_to))
@@ -271,9 +277,11 @@ def test_multiple_detectors_without_aggregation(self, input, expected):
             threshold_criterion=FixedThreshold(2),
             model_id="zscore_x2"))
 
-    with beam.Pipeline() as p:
+    with TestPipeline(
+        additional_pipeline_args=["--experiments=prism_disable_sdf_split"
+                                  ]) as p:
       result = (
-          p | beam.Create(input)
+          p | beam.Create(input, reshuffle=False)
           | AnomalyDetection(EnsembleAnomalyDetector(sub_detectors)))
 
       if isinstance(input[0], tuple):
@@ -298,9 +306,11 @@ def test_multiple_sub_detectors_with_aggregation(self, input, expected):
             threshold_criterion=FixedThreshold(2),
             model_id="zscore_x2"))
 
-    with beam.Pipeline() as p:
+    with TestPipeline(
+        additional_pipeline_args=["--experiments=prism_disable_sdf_split"
+                                  ]) as p:
       result = (
-          p | beam.Create(input)
+          p | beam.Create(input, reshuffle=False)
           | AnomalyDetection(
               EnsembleAnomalyDetector(
                   sub_detectors, aggregation_strategy=AnyVote())))
@@ -328,9 +338,11 @@ def test_ensemble_preserves_model_state(self, keyed):
                 for model_id in ('x1', 'x2')
                 for label in [-2, -2] + [0] * (count - 2)]
 
-    with TestPipeline() as p:
+    with TestPipeline(
+        additional_pipeline_args=["--experiments=prism_disable_sdf_split"
+                                  ]) as p:
       result = (
-          p | beam.Create(input)
+          p | beam.Create(input, reshuffle=False)
           | AnomalyDetection(EnsembleAnomalyDetector(detectors)))
       if not keyed:
         result = result | beam.WithKeys(0)
@@ -648,10 +660,12 @@ def test_dofn_on_single_prediction(self):
                 AnomalyResult(
                     R, [AnomalyPrediction(score=30, label=1, threshold=20)]))),
     ]
-    with TestPipeline() as p:
+    with TestPipeline(
+        additional_pipeline_args=["--experiments=prism_disable_sdf_split"
+                                  ]) as p:
       result = (
           p
-          | beam.Create(input)
+          | beam.Create(input, reshuffle=False)
           # use median just for test convenience
           | beam.ParDo(
               _StatefulThresholdDoFn(

sdks/python/apache_beam/yaml/integration_tests.py

@@ -107,6 +107,7 @@ def get_impl(self):
 from apache_beam.options.pipeline_options import GoogleCloudOptions
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.utils import python_callable
+from apache_beam.utils import retry
 from apache_beam.yaml import yaml_provider
 from apache_beam.yaml import yaml_transform
 from apache_beam.yaml.conftest import yaml_test_files_dir
@@ -411,6 +412,15 @@ def temp_mongodb_table():
     mongo_container.start()
     mongo_uri = mongo_container.get_connection_url()
 
+    # MongoDbContainer's entrypoint restarts mongod after init scripts; wait
+    # until the server is stably accepting connections and responds to ping.
+    @retry.with_exponential_backoff(num_retries=5, initial_delay_secs=1)
+    def _wait_for_mongo():
+      mongo_client = mongo_container.get_connection_client()
+      mongo_client.admin.command('ping')
+
+    _wait_for_mongo()
+
     db_name = f'db_{uuid.uuid4().hex}'
     collection_name = f'collection_{uuid.uuid4().hex}'
 

sdks/python/apache_beam/yaml/test_utils/datadog_test_utils.py

@@ -116,16 +116,14 @@ def temp_fake_datadog_server(expected_records=None):
       logging.error(
           "Error interacting with temporary fake Datadog server: %s", err)
       raise err
-    finally:
-      if expected_records is not None:
 
-        canonicalize = lambda rec: json.dumps(rec, sort_keys=True)
+  if expected_records is not None:
+    canonicalize = lambda rec: json.dumps(rec, sort_keys=True)
 
-        actual_strs = sorted([canonicalize(r) for r in received])
-        expected_strs = sorted([canonicalize(e) for e in expected_records])
+    actual_strs = sorted(set(canonicalize(r) for r in received))
+    expected_strs = sorted(set(canonicalize(e) for e in expected_records))
 
-        assert actual_strs == expected_strs, (
-            f"Mismatch in recorded Datadog events!\n"
-            f"Expected: {expected_strs}\n"
-            f"Actual:   {actual_strs}"
-        )
+    assert actual_strs == expected_strs, (
+        f"Mismatch in recorded Datadog events!\n"
+        f"Expected: {expected_strs}\n"
+        f"Actual:   {actual_strs}")