FrankYang0529 commented on code in PR #76218:
URL: https://github.com/apache/airflow/pull/76218#discussion_r4562274561
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -2177,6 +2208,81 @@ async def
test_create_triggers_emits_creation_duration_metric(
await runner.cleanup_finished_triggers()
+ @pytest.mark.asyncio
+ @patch("airflow.jobs.triggerer_job_runner._DRAIN_YIELD_INTERVAL", 0)
+ async def
test_cleanup_finished_triggers_visits_only_finished_triggers(self):
+ runner = TriggerRunner()
+ running_ids = set(range(1, 101))
+ runner.to_create.extend(_make_sleeping_trigger_workload(trigger_id)
for trigger_id in running_ids)
+ runner.to_create.append(
+ workloads.RunTrigger.model_construct(
+ id=101,
+ ti=None,
+
classpath=f"{SuccessTrigger.__module__}.{SuccessTrigger.__name__}",
+ encrypted_kwargs='{"__type":"dict", "__var":{}}',
+ )
+ )
+ await runner.create_triggers()
+ try:
+ await asyncio.wait([runner.triggers[101]["task"]])
+
+ finished_ids, yields = await
_await_counting_yields(runner.cleanup_finished_triggers())
+
+ assert finished_ids == [101]
+ assert set(runner.triggers) == running_ids
+ assert yields == 1
+ finally:
+ for details in runner.triggers.values():
+ details["task"].cancel()
+ await asyncio.gather(*(d["task"] for d in
runner.triggers.values()), return_exceptions=True)
+
+ @pytest.mark.asyncio
+ @patch("airflow.jobs.triggerer_job_runner._maybe_yield", autospec=True)
+ async def
test_cleanup_finished_triggers_handles_only_ids_queued_when_it_starts(self,
mock_maybe_yield):
+ runner = TriggerRunner()
+ ids_finishing_during_yield: list[int] = []
+
+ def finish_triggers_during_yield(last_yield: float) -> float:
+ runner._finished_trigger_ids.extend(ids_finishing_during_yield)
+ ids_finishing_during_yield.clear()
+ return last_yield
+
+ mock_maybe_yield.side_effect = finish_triggers_during_yield
+ runner.to_create.extend(_make_sleeping_trigger_workload(trigger_id)
for trigger_id in (1, 2))
+ await runner.create_triggers()
+ tasks = [details["task"] for details in runner.triggers.values()]
+ for task in tasks:
+ task.cancel()
+ await asyncio.wait(tasks)
+ runner._finished_trigger_ids.remove(2)
+ ids_finishing_during_yield.append(2)
+
+ assert await runner.cleanup_finished_triggers() == [1]
+ assert set(runner.triggers) == {2}
+ assert await runner.cleanup_finished_triggers() == [2]
+ assert not runner.triggers
+
+ @pytest.mark.asyncio
+ @patch("airflow.jobs.triggerer_job_runner._DRAIN_YIELD_INTERVAL", 3600)
Review Comment:
With the interval at 3600, `_maybe_yield()` always returns without yielding.
As a result, this test gets `(0, 0, 0)` whether or not the loops call
`_maybe_yield()`.
Could this test also run with the interval at 0? With the interval at 0,
every `_maybe_yield()` call yields once.
```diff
@pytest.mark.asyncio
- @patch("airflow.jobs.triggerer_job_runner._DRAIN_YIELD_INTERVAL", 3600)
- async def
test_create_cancel_and_cleanup_yield_by_time_not_per_trigger(self):
+ @pytest.mark.parametrize(
+ ("interval", "expected_yields"),
+ [
+ pytest.param(3600, (0, 0, 0), id="interval-not-reached"),
+ pytest.param(0, (200, 100, 100), id="every-call-yields"),
+ ],
+ )
+ async def test_create_cancel_and_cleanup_yield_by_time_not_per_trigger(
+ self, monkeypatch, interval, expected_yields
+ ):
+
monkeypatch.setattr("airflow.jobs.triggerer_job_runner._DRAIN_YIELD_INTERVAL",
interval)
runner = TriggerRunner()
@@
- assert (create_yields, cancel_yields, cleanup_yields) == (0, 0, 0)
+ assert (create_yields, cancel_yields, cleanup_yields) ==
expected_yields
```
--
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]
