Skip to content

Commit

Permalink
Added a regression test for #15952
Browse files Browse the repository at this point in the history
  • Loading branch information
jdjfisher committed Nov 15, 2024
1 parent c4dad5b commit 75d87f4
Showing 1 changed file with 29 additions and 1 deletion.
30 changes: 29 additions & 1 deletion tests/workers/test_base_worker.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
):
Expand Down

0 comments on commit 75d87f4

Please sign in to comment.