11import asyncio
2- import time
32
43import pytest
54
@@ -16,18 +15,18 @@ async def test_wait_result() -> None:
1615 runs = 0
1716
1817 @broker .task (retry_on_error = True )
19- def run_task () -> str :
18+ async def run_task () -> str :
2019 nonlocal runs
2120
2221 if runs == 0 :
2322 runs += 1
2423 raise Exception ("Retry" )
2524
26- time .sleep (0.2 )
2725 return "hello world!"
2826
2927 task = await run_task .kiq ()
3028 resp = await task .wait_result (0.1 , timeout = 1 )
29+ assert runs == 1
3130
3231 assert resp .return_value == "hello world!"
3332
@@ -39,24 +38,28 @@ async def test_wait_result_error() -> None:
3938 SimpleRetryMiddleware (no_result_on_retry = False ),
4039 )
4140 runs = 0
41+ lock = asyncio .Lock ()
4242
4343 @broker .task (retry_on_error = True )
44- def run_task () -> str :
45- nonlocal runs
44+ async def run_task () -> str :
45+ nonlocal runs , lock
46+
47+ await lock .acquire ()
4648
4749 if runs == 0 :
4850 runs += 1
4951 raise ValueError ("Retry" )
5052
51- time .sleep (0.2 )
5253 return "hello world!"
5354
5455 task = await run_task .kiq ()
5556 resp = await task .wait_result (0.1 , timeout = 1 )
56- with pytest .raises (ValueError ):
57- resp .raise_for_error ()
57+ assert resp .is_err
58+ assert runs == 1
59+
60+ broker .result_backend .results .pop (task .task_id ) # type: ignore
61+ lock .release ()
5862
59- await asyncio .sleep (0.2 )
6063 resp = await task .wait_result (timeout = 1 )
6164 assert resp .return_value == "hello world!"
6265
@@ -67,32 +70,34 @@ async def test_wait_result_no_result() -> None:
6770 broker = InMemoryBroker ().with_middlewares (
6871 SimpleRetryMiddleware (no_result_on_retry = False ),
6972 )
70- done = False
73+ done = asyncio . Event ()
7174 runs = 0
75+ lock = asyncio .Lock ()
7276
7377 @broker .task (retry_on_error = True )
74- def run_task () -> str :
75- nonlocal runs , done
78+ async def run_task () -> str :
79+ nonlocal runs , done , lock
80+
81+ await lock .acquire ()
7682
7783 if runs == 0 :
7884 runs += 1
7985 raise ValueError ("Retry" )
8086
81- time .sleep (0.2 )
82- done = True
87+ done .set ()
8388 raise NoResultError
8489
8590 task = await run_task .kiq ()
8691 resp = await task .wait_result (0.1 , timeout = 1 )
8792 with pytest .raises (ValueError ):
8893 resp .raise_for_error ()
8994
90- await asyncio .sleep (0.2 )
91- resp = await task .wait_result (timeout = 1 )
92- with pytest .raises (ValueError ):
93- resp .raise_for_error ()
95+ broker .result_backend .results .pop (task .task_id ) # type: ignore
96+ lock .release ()
9497
95- assert done
98+ assert await asyncio .wait_for (done .wait (), timeout = 1 )
99+ with pytest .raises (KeyError ):
100+ await broker .result_backend .get_result (task .task_id )
96101
97102
98103@pytest .mark .anyio
0 commit comments