pingzh commented on code in PR #21877:
URL: https://github.com/apache/airflow/pull/21877#discussion_r851394252


##########
tests/jobs/test_local_task_job.py:
##########
@@ -276,39 +278,26 @@ def test_heartbeat_failed_fast(self):
                 delta = (time2 - time1).total_seconds()
                 assert abs(delta - job.heartrate) < 0.5
 
-    @patch('airflow.utils.process_utils.subprocess.check_call')
-    @patch.object(StandardTaskRunner, 'return_code')
-    def test_mark_success_no_kill(self, mock_return_code, _check_call, caplog, 
dag_maker):
+    def test_mark_success_no_kill(self, caplog, get_test_dag):
         """
         Test that ensures that mark_success in the UI doesn't cause
         the task to fail, and that the task exits
         """
-        session = settings.Session()
-
-        def task_function(ti):
-            assert ti.state == State.RUNNING
-            # Simulate marking this successful in the UI
-            ti.state = State.SUCCESS
-            session.merge(ti)
-            session.commit()
-            # The below code will not run as heartbeat will detect change of 
state
-            time.sleep(10)
-
-        with dag_maker('test_mark_success'):
-            task = PythonOperator(task_id="task1", 
python_callable=task_function)
-        dr = dag_maker.create_dagrun()
+        dag = get_test_dag('test_mark_state')
+        with create_session() as session:

Review Comment:
   good call. thanks. updated.



-- 
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]

Reply via email to