1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
|
import pytest
import asyncio
from aiostream.test_utils import event_loop
from aiostream.aiter_utils import AsyncIteratorContext, aitercontext, anext
# Pytest fixtures
event_loop
# Some async iterators for testing
async def agen():
for x in range(5):
await asyncio.sleep(1)
yield x
async def silence_agen():
while True:
try:
yield 1
except BaseException:
return
async def reraise_agen():
while True:
try:
yield 1
except BaseException as exc:
raise RuntimeError from exc
async def stuck_agen():
try:
yield 1
except BaseException:
yield 2
class not_an_agen(list):
def __aiter__(self):
return self
async def __anext__(self):
try:
return self.pop()
except IndexError:
raise StopAsyncIteration
# Tests
@pytest.mark.asyncio
async def test_simple_aitercontext(event_loop):
async with aitercontext(agen()) as safe_gen:
# Cannot enter twice
with pytest.raises(RuntimeError):
async with safe_gen:
pass
it = iter(range(5))
async for item in safe_gen:
assert item == next(it)
assert event_loop.steps == [1] * 5
# Exiting is idempotent
await safe_gen.__aexit__(None, None, None)
await safe_gen.aclose()
with pytest.raises(RuntimeError):
await anext(safe_gen)
with pytest.raises(RuntimeError):
async with safe_gen:
pass
with pytest.raises(RuntimeError):
await safe_gen.athrow(ValueError())
@pytest.mark.asyncio
async def test_athrow_in_aitercontext(event_loop):
async with aitercontext(agen()) as safe_gen:
assert await safe_gen.__anext__() == 0
with pytest.raises(ZeroDivisionError):
await safe_gen.athrow(ZeroDivisionError())
async for _ in safe_gen:
assert False # No more items
@pytest.mark.asyncio
async def test_aitercontext_wrong_usage(event_loop):
safe_gen = aitercontext(agen())
with pytest.warns(UserWarning):
await anext(safe_gen)
with pytest.raises(TypeError):
AsyncIteratorContext(None)
with pytest.raises(TypeError):
AsyncIteratorContext(safe_gen)
@pytest.mark.asyncio
async def test_raise_in_aitercontext(event_loop):
with pytest.raises(ZeroDivisionError):
async with aitercontext(agen()) as safe_gen:
async for _ in safe_gen:
1 / 0
with pytest.raises(ZeroDivisionError):
async with aitercontext(agen()) as safe_gen:
async for _ in safe_gen:
pass
1 / 0
with pytest.raises(GeneratorExit):
async with aitercontext(agen()) as safe_gen:
async for _ in safe_gen:
raise GeneratorExit
with pytest.raises(GeneratorExit):
async with aitercontext(agen()) as safe_gen:
async for _ in safe_gen:
pass
raise GeneratorExit
@pytest.mark.asyncio
async def test_silence_exception_in_aitercontext(event_loop):
async with aitercontext(silence_agen()) as safe_gen:
async for item in safe_gen:
assert item == 1
1 / 0
# Silencing a generator exit is forbidden
with pytest.raises(GeneratorExit):
async with aitercontext(silence_agen()) as safe_gen:
async for _ in safe_gen:
raise GeneratorExit
@pytest.mark.asyncio
async def test_reraise_exception_in_aitercontext(event_loop):
with pytest.raises(RuntimeError) as info:
async with aitercontext(reraise_agen()) as safe_gen:
async for item in safe_gen:
assert item == 1
1 / 0
assert type(info.value.__cause__) is ZeroDivisionError
with pytest.raises(RuntimeError) as info:
async with aitercontext(reraise_agen()) as safe_gen:
async for item in safe_gen:
assert item == 1
raise GeneratorExit
assert type(info.value.__cause__) is GeneratorExit
@pytest.mark.asyncio
async def test_stuck_in_aitercontext(event_loop):
with pytest.raises(RuntimeError) as info:
async with aitercontext(stuck_agen()) as safe_gen:
async for item in safe_gen:
assert item == 1
1 / 0
assert "didn't stop after athrow" in str(info.value)
with pytest.raises(RuntimeError) as info:
async with aitercontext(stuck_agen()) as safe_gen:
async for item in safe_gen:
assert item == 1
raise GeneratorExit
# GeneratorExit relies on aclose, not athrow
# The message is a bit different
assert "async generator ignored GeneratorExit" in str(info.value)
@pytest.mark.asyncio
async def test_not_an_agen_in_aitercontext(event_loop):
async with aitercontext(not_an_agen([1])) as safe_gen:
async for item in safe_gen:
assert item == 1
with pytest.raises(ZeroDivisionError):
async with aitercontext(not_an_agen([1])) as safe_gen:
async for item in safe_gen:
1 / 0
|
