返回提交历史
Added
etc/unittest/test_provider_asyncio.py
+31
-0
Modified
g4f/providers/base_provider.py
+27
-19
XFEstudio/gpt4free
Refactor wait_for function to improve timeout handling and cleanup logic
c1d57511
代码差异
2 个文件
+58
-19
@@ -0,0 +1,31 @@
1
"""Tests for async provider timeout cleanup."""
2
3
import asyncio
4
import unittest
5
6
from g4f.providers.base_provider import wait_for
7
8
9
class TestWaitFor(unittest.TestCase):
10
def test_timeout_closes_wrapped_generator(self):
11
closed = False
12
13
async def response():
14
nonlocal closed
15
try:
16
await asyncio.Event().wait()
17
yield "unreachable"
18
finally:
19
closed = True
20
21
async def run():
22
with self.assertRaises(TimeoutError):
23
async for _ in wait_for(response(), timeout=0.01):
24
pass
25
self.assertTrue(closed)
26
27
asyncio.run(run())
28
29
30
if __name__ == "__main__":
31
unittest.main()
@@ -105,26 +105,34 @@ PARAMETER_EXAMPLES = {
105
105
106
106
107
107
async def wait_for(response: AsyncIterator, timeout: int = None) -> AsyncIterator:
108
if timeout is not None:
109
while True:
108
try:
109
if timeout is not None:
110
while True:
111
try:
112
async def wait_for_next():
113
try:
114
return await response.__anext__()
115
except TimeoutError as e:
116
raise TimeoutError(str(e) or "The operation timed out") from e
117
yield await asyncio.wait_for(wait_for_next(), timeout=timeout)
118
except TimeoutError as e:
119
if str(e):
120
raise TimeoutError(str(e)) from e
121
raise TimeoutError(
122
"The operation timed out after {} seconds".format(timeout)
123
) from e
124
except StopAsyncIteration:
125
break
126
else:
127
async for chunk in response:
128
yield chunk
129
finally:
130
close = getattr(response, "aclose", None)
131
if close is not None:
110
132
try:
111
async def wait_for_next():
112
try:
113
return await response.__anext__()
114
except TimeoutError as e:
115
raise TimeoutError(str(e) or "The operation timed out") from e
116
yield await asyncio.wait_for(wait_for_next(), timeout=timeout)
117
except TimeoutError as e:
118
if str(e):
119
raise TimeoutError(str(e)) from e
120
raise TimeoutError(
121
"The operation timed out after {} seconds".format(timeout)
122
) from e
123
except StopAsyncIteration:
124
break
125
else:
126
async for chunk in response:
127
yield chunk
133
await close()
134
except Exception:
135
pass
128
136
129
137
130
138
def get_async_provider_method(provider: type) -> Optional[callable]: