[GitHub] [airflow] uranusjr commented on a change in pull request #17581: Remove the use of multiprocessing in TestLocalTaskJob and Improve Tests

[GitBox <gi...@apache.org>]

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-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org