dabla commented on code in PR #37937:
URL: https://github.com/apache/airflow/pull/37937#discussion_r1516017667
##########
tests/models/test_baseoperatormeta.py:
##########
@@ -0,0 +1,64 @@
+import os
+from typing import Any
+from unittest import TestCase
+from unittest.mock import Mock, patch, MagicMock
+
+import pendulum
+from sqlalchemy.orm import Session
+
+from airflow import AirflowException, DAG
+from airflow.models import TaskInstance, DagRun
+from airflow.models.baseoperator import BaseOperator
+from airflow.utils.context import Context
+from airflow.utils.state import DagRunState
+
+
+class HelloWorldOperator(BaseOperator):
+ called = False
+
+ def execute(self, context: Context) -> Any:
+ HelloWorldOperator.called = True
+ return f"Hello {self.owner}!"
+
+
+class IterableSession(Session):
+ def __next__(self):
+ pass
+
+
+class BaseOperatorMetaTestCase(TestCase):
+
+ @patch.dict(os.environ, {'AIRFLOW__CORE__UNIT_TEST_MODE': 'False'})
+ def test_executor_safeguard_when_unauthorized(self):
+ with self.assertRaises(AirflowException):
+ dag = DAG(dag_id="hello_world")
+ context = MagicMock(spec=Context)
+
+ HelloWorldOperator(task_id="task_id",
dag=dag).execute(context=context)
+
+ @patch("sqlalchemy.orm.Session.__init__")
+ @patch.dict(os.environ, {'AIRFLOW__CORE__UNIT_TEST_MODE': 'False'})
+ def test_executor_safeguard_when_authorized(self, mock_session: MagicMock):
+ session = MagicMock(spec=IterableSession)
+ mock_session.return_value = session
+ session.__iter__.return_value = iter({})
+
+ dag = DAG(dag_id="hello_world")
+ context = {"params": {}}
Review Comment:
I think I need to add another test case, namely a test case where a
PythonOperator is being executed through the TaskInstance but has a callable
which manually calls the execute of the HelloWorldOperator, as this is what
will happen in real on Airflow and that should raise an AirflowException.
--
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]
