uranusjr commented on a change in pull request #17581:
URL: https://github.com/apache/airflow/pull/17581#discussion_r690217140
##########
File path: tests/jobs/test_local_task_job.py
##########
@@ -372,16 +361,17 @@ def test_localtaskjob_maintain_heart_rate(self):
# loop in _execute()
return_codes = [None, 0]
- def multi_return_code():
- return return_codes.pop(0)
+ def multi_return_code(*args, **kwargs):
+ while return_codes:
+ return return_codes.pop(0)
time_start = time.time()
- with patch.object(StandardTaskRunner, 'start', return_value=None) as
mock_start:
- with patch.object(StandardTaskRunner, 'return_code') as
mock_ret_code:
- mock_ret_code.side_effect = multi_return_code
- job1.run()
- assert mock_start.call_count == 1
- assert mock_ret_code.call_count == 2
+
+ mock_return_code.side_effect = multi_return_code
Review comment:
```suggestion
mock_return_code.side_effect = [None, 0]
```
[Documentation on
this](https://docs.python.org/3/library/unittest.mock.html#quick-guide)
--
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: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
