This is an automated email from the ASF dual-hosted git repository.
reuvenlax pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 829f41d0b49 Fix flaky Python tests (#40297)
829f41d0b49 is described below
commit 829f41d0b4976ebda2a2e25e6866dfe8c08d7f03
Author: Reuven Lax <[email protected]>
AuthorDate: Wed Oct 14 20:40:23 2026 -0700
Fix flaky Python tests (#40297)
---
.../examples/cookbook/bigtableio_it_test.py | 20 +++++++++---
.../examples/streaming_wordcount_it_test.py | 11 +++++--
sdks/python/apache_beam/io/gcp/bigtableio.py | 5 +++
sdks/python/apache_beam/io/gcp/bigtableio_test.py | 17 +++++++---
.../apache_beam/io/gcp/pubsub_integration_test.py | 20 ++++++++----
.../apache_beam/ml/anomaly/transforms_test.py | 38 +++++++++++++++-------
sdks/python/apache_beam/yaml/integration_tests.py | 10 ++++++
.../yaml/test_utils/datadog_test_utils.py | 18 +++++-----
8 files changed, 99 insertions(+), 40 deletions(-)
diff --git a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
index 84087dee560..bca6881a714 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
@@ -35,6 +35,7 @@ from apache_beam.metrics.metric import MetricsFilter
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 @@ class BigtableIOWriteTest(unittest.TestCase):
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)
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
index 9ed27a500a7..c67e8e90e5a 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
@@ -66,12 +66,19 @@ class StreamingWordCountIT(unittest.TestCase):
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 @@ class StreamingWordCountIT(unittest.TestCase):
# 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,
diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py
b/sdks/python/apache_beam/io/gcp/bigtableio.py
index b8d21242806..d335fd9fc5d 100644
--- a/sdks/python/apache_beam/io/gcp/bigtableio.py
+++ b/sdks/python/apache_beam/io/gcp/bigtableio.py
@@ -290,6 +290,11 @@ class WriteToBigTable(beam.PTransform):
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',
diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py
b/sdks/python/apache_beam/io/gcp/bigtableio_test.py
index 4aba9a4062f..16fd41e6d1d 100644
--- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py
+++ b/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
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 89fd4461beb..d665f64168e 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py
@@ -36,6 +36,7 @@ from apache_beam.runners.runner import PipelineState
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 @@ class PubSubIntegrationTest(unittest.TestCase):
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)
diff --git a/sdks/python/apache_beam/ml/anomaly/transforms_test.py
b/sdks/python/apache_beam/ml/anomaly/transforms_test.py
index 72c16841d5f..4da9d6c9178 100644
--- a/sdks/python/apache_beam/ml/anomaly/transforms_test.py
+++ b/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 TestAnomalyDetection(unittest.TestCase):
])
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 @@ class TestAnomalyDetection(unittest.TestCase):
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 @@ class TestAnomalyDetection(unittest.TestCase):
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 @@ class TestAnomalyDetection(unittest.TestCase):
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 @@ class TestStatefulThresholdDoFn(unittest.TestCase):
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(
diff --git a/sdks/python/apache_beam/yaml/integration_tests.py
b/sdks/python/apache_beam/yaml/integration_tests.py
index 107d55df19b..9af075ef5d8 100644
--- a/sdks/python/apache_beam/yaml/integration_tests.py
+++ b/sdks/python/apache_beam/yaml/integration_tests.py
@@ -107,6 +107,7 @@ from apache_beam.io.gcp.spanner_wrapper import
SpannerWrapper
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}'
diff --git a/sdks/python/apache_beam/yaml/test_utils/datadog_test_utils.py
b/sdks/python/apache_beam/yaml/test_utils/datadog_test_utils.py
index 4a87a346668..f947a472005 100644
--- a/sdks/python/apache_beam/yaml/test_utils/datadog_test_utils.py
+++ b/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}")
