[GitHub] [airflow] ashb commented on a change in pull request #16852: Fix impersonation issue with LocalTaskJob

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

ashb commented on a change in pull request #16852:
URL: https://github.com/apache/airflow/pull/16852#discussion_r665258586



##########
File path: tests/jobs/test_local_task_job.py
##########
@@ -147,6 +147,52 @@ def test_localtaskjob_heartbeat(self):
         with pytest.raises(AirflowException):
             job1.heartbeat_callback()
 
+    @mock.patch('airflow.jobs.local_task_job.psutil')
+    def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock):
+        session = settings.Session()
+        dag = DAG('test_localtaskjob_heartbeat', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'})
+
+        with dag:
+            op1 = DummyOperator(task_id='op1', run_as_user='myuser')
+
+        dag.clear()
+        dr = dag.create_dagrun(
+            run_id="test",
+            state=State.SUCCESS,
+            execution_date=DEFAULT_DATE,
+            start_date=DEFAULT_DATE,
+            session=session,
+        )
+
+        ti = dr.get_task_instance(task_id=op1.task_id, session=session)
+        ti.state = State.RUNNING
+        ti.pid = 2
+        ti.hostname = get_hostname()
+        session.commit()
+
+        job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor())
+        ti.task = op1
+        ti.refresh_from_task(op1)
+        job1.task_runner = StandardTaskRunner(job1)
+        job1.task_runner.process = mock.Mock()
+        job1.task_runner.process.pid = 2
+        with pytest.raises(AirflowException):
+            job1.heartbeat_callback()

Review comment:
       What exception is this raising? We should use the "match" parameter to assert _something_ about what the error actually is.

##########
File path: tests/jobs/test_local_task_job.py
##########
@@ -147,6 +147,52 @@ def test_localtaskjob_heartbeat(self):
         with pytest.raises(AirflowException):
             job1.heartbeat_callback()
 
+    @mock.patch('airflow.jobs.local_task_job.psutil')
+    def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock):
+        session = settings.Session()
+        dag = DAG('test_localtaskjob_heartbeat', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'})
+
+        with dag:
+            op1 = DummyOperator(task_id='op1', run_as_user='myuser')
+
+        dag.clear()
+        dr = dag.create_dagrun(
+            run_id="test",
+            state=State.SUCCESS,
+            execution_date=DEFAULT_DATE,
+            start_date=DEFAULT_DATE,
+            session=session,
+        )
+
+        ti = dr.get_task_instance(task_id=op1.task_id, session=session)
+        ti.state = State.RUNNING
+        ti.pid = 2
+        ti.hostname = get_hostname()
+        session.commit()
+
+        job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor())
+        ti.task = op1
+        ti.refresh_from_task(op1)
+        job1.task_runner = StandardTaskRunner(job1)
+        job1.task_runner.process = mock.Mock()
+        job1.task_runner.process.pid = 2
+        with pytest.raises(AirflowException):
+            job1.heartbeat_callback()
+
+        job1.task_runner.process.pid = 1
+        psutil_mock.Process.return_value.ppid.return_value = 1
+        ti.state = State.RUNNING
+        ti.pid = 2
+        assert ti.run_as_user
+        session.merge(ti)
+        session.commit()
+        assert ti.pid != os.getpid()

Review comment:
       ```suggestion
   ```
   
   I don't think this is necessary.




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