kacpermuda commented on code in PR #70760:
URL: https://github.com/apache/airflow/pull/70760#discussion_r3703919343
##########
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:
Correct, used `testing_team` and `testing_dag_bundle` fixtures instead, and
linked them together in a new small fixture that I reused in all tests. Should
be clean now, thanks for spotting that.
--
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]
