96 lines
2.5 KiB
Python
96 lines
2.5 KiB
Python
import unittest
|
|
|
|
from pwo import retry, async_retry, async_test, AsyncQueueIterator, aenumerate
|
|
from asyncio import Queue
|
|
|
|
|
|
class PrivateTest(unittest.TestCase):
|
|
def test_retry_until_success(self):
|
|
max_attempts = 20
|
|
attempt = 0
|
|
|
|
expected_result = object()
|
|
|
|
@retry(max_attempts=max_attempts, initial_delay=0)
|
|
def foo():
|
|
nonlocal attempt
|
|
attempt += 1
|
|
if attempt < 10:
|
|
raise Exception()
|
|
else:
|
|
return expected_result
|
|
|
|
self.assertEqual(expected_result, foo())
|
|
self.assertEqual(10, attempt)
|
|
|
|
def test_retry_until_max_attempt(self):
|
|
max_attempts = 20
|
|
attempt = 0
|
|
|
|
@retry(max_attempts=max_attempts, initial_delay=0)
|
|
def bar():
|
|
nonlocal attempt
|
|
attempt += 1
|
|
raise Exception()
|
|
|
|
with self.assertRaises(Exception):
|
|
bar()
|
|
self.assertEqual(max_attempts, attempt)
|
|
|
|
@async_test
|
|
async def test_async_retry_until_success(self):
|
|
max_attempts = 20
|
|
attempt = 0
|
|
|
|
expected_result = object()
|
|
|
|
@async_retry(max_attempts=max_attempts, initial_delay=0)
|
|
async def foo():
|
|
nonlocal attempt
|
|
attempt += 1
|
|
if attempt < 10:
|
|
raise Exception()
|
|
else:
|
|
return expected_result
|
|
|
|
self.assertEqual(expected_result, await foo())
|
|
self.assertEqual(10, attempt)
|
|
|
|
@async_test
|
|
async def test_async_retry_until_max_attempt(self):
|
|
max_attempts = 20
|
|
attempt = 0
|
|
|
|
@async_retry(max_attempts=max_attempts, initial_delay=0)
|
|
async def bar():
|
|
nonlocal attempt
|
|
attempt += 1
|
|
raise Exception()
|
|
|
|
with self.assertRaises(Exception):
|
|
await bar()
|
|
self.assertEqual(max_attempts, attempt)
|
|
|
|
@async_test
|
|
async def test_async_queue_iterator(self):
|
|
queue = Queue()
|
|
queue_size = 10
|
|
objects = [object() for _ in range(queue_size)]
|
|
|
|
async def poll() -> int:
|
|
completed = 0
|
|
async for i, obj in aenumerate(AsyncQueueIterator(queue)):
|
|
self.assertIs(objects[i], obj)
|
|
completed += 1
|
|
return completed
|
|
|
|
handle = poll()
|
|
|
|
for o in objects:
|
|
queue.put_nowait(o)
|
|
queue.put_nowait(None)
|
|
processed = await handle
|
|
self.assertEqual(queue_size, processed)
|
|
|
|
|