From 6a12d890eea72e738802679aa39ad6b636e771eb Mon Sep 17 00:00:00 2001 From: bvolpato Date: Sun, 25 Oct 2026 02:02:52 -0400 Subject: [PATCH] [Python] Make remote inference retry tests deterministic --- sdks/python/apache_beam/ml/inference/base_test.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index dec1c2906cd2..1a6d7f669df2 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2415,7 +2415,7 @@ def create_client(self): return FakeModel() def request(self, batch, model, inference_args=None) -> Iterable[int]: - raise Exception + raise RuntimeError('fake remote inference request failed') def batch_elements_kwargs(self): return { @@ -2475,22 +2475,29 @@ def test_normal_model_execution(self): actual = pcoll | base.RunInference(FakeRemoteModelHandler()) assert_that(actual, equal_to(expected), label='assert:inferences') - def test_repeated_requests_fail(self): + @unittest.mock.patch.object(base.ReactiveThrottler, 'throttle') + @unittest.mock.patch.object(base.retry.Clock, 'sleep') + def test_repeated_requests_fail(self, sleep, unused_throttle): test_pipeline = TestPipeline() - with self.assertRaises(Exception): + with self.assertRaisesRegex(RuntimeError, + 'fake remote inference request failed'): _ = ( test_pipeline | beam.Create([1, 2, 3, 4]) | base.RunInference(FakeAlwaysFailsRemoteModelHandler())) test_pipeline.run() + self.assertEqual(sleep.call_count, 2) - def test_works_on_retry(self): + @unittest.mock.patch.object(base.ReactiveThrottler, 'throttle') + @unittest.mock.patch.object(base.retry.Clock, 'sleep') + def test_works_on_retry(self, sleep, unused_throttle): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] expected = [example + 1 for example in examples] pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(FakeFailsOnceRemoteModelHandler()) assert_that(actual, equal_to(expected), label='assert:inferences') + sleep.assert_called() def test_exception_on_load_model_override(self): with self.assertRaises(Exception):