gemini-code-assist[bot] commented on code in PR #37581:
URL: https://github.com/apache/beam/pull/37581#discussion_r2800812126
##########
sdks/python/apache_beam/runners/dataflow/dataflow_runner_test.py:
##########
@@ -195,6 +195,22 @@ def get_job_side_effect(*args, **kwargs):
result = duration_timedout_result.wait_until_finish(5000)
self.assertEqual(result, PipelineState.RUNNING)
+ with mock.patch('time.time', mock.MagicMock(side_effect=[1, 9, 9, 20,
20])):
+ duration_timedout_runner = MockDataflowRunner(
+ [values_enum.JOB_STATE_PAUSING])
+ duration_timedout_result = DataflowPipelineResult(
+ duration_timedout_runner.job, duration_timedout_runner, options)
+ result = duration_timedout_result.wait_until_finish(5000)
+ self.assertEqual(result, PipelineState.PAUSING)
+
+ with mock.patch('time.time', mock.MagicMock(side_effect=[1, 9, 9, 20,
20])):
+ duration_timedout_runner = MockDataflowRunner(
+ [values_enum.JOB_STATE_PAUSED])
+ duration_timedout_result = DataflowPipelineResult(
+ duration_timedout_runner.job, duration_timedout_runner, options)
+ result = duration_timedout_result.wait_until_finish(5000)
+ self.assertEqual(result, PipelineState.PAUSED)
Review Comment:
These two test blocks are very similar. To avoid code duplication and
improve maintainability, consider refactoring them into a single parameterized
test using a loop and `self.subTest`.
```suggestion
for job_state, pipeline_state in [
(values_enum.JOB_STATE_PAUSING, PipelineState.PAUSING),
(values_enum.JOB_STATE_PAUSED, PipelineState.PAUSED),
]:
with self.subTest(job_state=job_state):
with mock.patch('time.time', mock.MagicMock(side_effect=[1, 9, 9,
20, 20])):
duration_timedout_runner = MockDataflowRunner([job_state])
duration_timedout_result = DataflowPipelineResult(
duration_timedout_runner.job, duration_timedout_runner,
options)
result = duration_timedout_result.wait_until_finish(5000)
self.assertEqual(result, pipeline_state)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
