ferruzzi commented on code in PR #70760:
URL: https://github.com/apache/airflow/pull/70760#discussion_r3693875203
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -9635,6 +9635,205 @@ def test_dag_timeout_notifies_with_timed_out_msg(self,
mock_get_listener_manager
assert call_args.kwargs["msg"] == "timed_out"
assert call_args.kwargs["dag_run"] == dag_run
+ @conf_vars({("core", "multi_team"): "true"})
+ @mock.patch("airflow.models.dagrun.get_listener_manager")
+ def test_dag_start_notifies_listener_with_team_name(self,
mock_get_listener_manager, dag_maker, session):
+ """Test that on_dag_run_running receives dag_run with _team_name
set."""
+ mock_listener_manager = MagicMock()
+ mock_get_listener_manager.return_value = mock_listener_manager
+
+ clear_db_teams()
+ clear_db_dag_bundles()
+
+ team = Team(name="test_team")
+ session.add(team)
+ session.flush()
+
+ bundle = DagBundleModel(name="test_bundle")
+ bundle.teams.append(team)
+ session.add(bundle)
+ session.flush()
+
+ with dag_maker(dag_id="test_dag_start_team",
bundle_name="test_bundle", session=session):
+ EmptyOperator(task_id="test_task")
+
+ dag_maker.create_dagrun(run_id="test_run", state=DagRunState.QUEUED)
+ session.commit()
+
+ mock_executor = MagicMock()
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[mock_executor])
+
+ self.job_runner._start_queued_dagruns(session)
+
+ mock_listener_manager.hook.on_dag_run_running.assert_called_once()
+ call_args = mock_listener_manager.hook.on_dag_run_running.call_args
+ assert call_args.kwargs["dag_run"]._team_name == "test_team"
+
+ @conf_vars({("core", "multi_team"): "true"})
+ @time_machine.travel(DEFAULT_DATE, tick=False)
+ @mock.patch("airflow.models.dagrun.get_listener_manager")
+ def test_dag_timeout_notifies_listener_with_team_name(
+ self, mock_get_listener_manager, dag_maker, session
+ ):
+ """Test that on_dag_run_failed receives dag_run with _team_name set
when a DAG times out."""
+ mock_listener_manager = MagicMock()
+ mock_get_listener_manager.return_value = mock_listener_manager
+
+ clear_db_teams()
+ clear_db_dag_bundles()
+
+ team = Team(name="test_team")
+ session.add(team)
+ session.flush()
+
+ bundle = DagBundleModel(name="test_bundle")
+ bundle.teams.append(team)
+ session.add(bundle)
+ session.flush()
Review Comment:
These tests all have a pretty big repetitive block at the top, maybe move
the setup into a shared fixture if none of the existing fixtures do what you
need? look at the `testing_team` fixture though, I think that covers most/all
of this boilerplate.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -9635,6 +9635,205 @@ def test_dag_timeout_notifies_with_timed_out_msg(self,
mock_get_listener_manager
assert call_args.kwargs["msg"] == "timed_out"
assert call_args.kwargs["dag_run"] == dag_run
+ @conf_vars({("core", "multi_team"): "true"})
+ @mock.patch("airflow.models.dagrun.get_listener_manager")
+ def test_dag_start_notifies_listener_with_team_name(self,
mock_get_listener_manager, dag_maker, session):
+ """Test that on_dag_run_running receives dag_run with _team_name
set."""
+ mock_listener_manager = MagicMock()
+ mock_get_listener_manager.return_value = mock_listener_manager
+
+ clear_db_teams()
+ clear_db_dag_bundles()
Review Comment:
Clearing the db at the beginning of the tests should be a code smell. I
believe this issue was fixed in #69093, if not, then we need to figure out what
tests are not cleaning up after themselves and fix them.
--
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]
