diff --git a/tests/workers/test_base_worker.py b/tests/workers/test_base_worker.py index 9839e0cf3029..03973598fe95 100644 --- a/tests/workers/test_base_worker.py +++ b/tests/workers/test_base_worker.py @@ -1,7 +1,7 @@ import uuid from typing import Any, Dict, Optional, Type from unittest import mock -from unittest.mock import MagicMock +from unittest.mock import MagicMock, Mock import httpx import pendulum @@ -338,6 +338,34 @@ def create_run_with_deployment_2(state): assert {flow_run.id for flow_run in submitted_flow_runs} == set(flow_run_ids[1:2]) +async def test_worker_releases_limit_slot_when_aborting_a_change_to_pending( + prefect_client: PrefectClient, worker_deployment_wq1, work_pool +): + """Regression test for https://github.com/PrefectHQ/prefect/issues/15952""" + + def create_run_with_deployment(state): + return prefect_client.create_flow_run_from_deployment( + worker_deployment_wq1.id, state=state + ) + + flow_run = await create_run_with_deployment( + Scheduled(scheduled_time=pendulum.now("utc").subtract(days=1)) + ) + + run_mock = AsyncMock() + release_mock = Mock() + + async with WorkerTestImpl(work_pool_name=work_pool.name, limit=1) as worker: + worker.run = run_mock + worker._propose_pending_state = AsyncMock(return_value=False) + worker._release_limit_slot = release_mock + + await worker.get_and_submit_flow_runs() + + run_mock.assert_not_called() + release_mock.assert_called_once_with(flow_run.id) + + async def test_worker_with_work_pool_and_limit( prefect_client: PrefectClient, worker_deployment_wq1, work_pool ):