diff --git a/goblin/gremlin_python_driver/driver.py b/goblin/gremlin_python_driver/driver.py index 33a723df55103a38cb783451538af61d86c8a0b9..dc8b206dd65612cfe6b4cfd4f01d980ca4c822d8 100644 --- a/goblin/gremlin_python_driver/driver.py +++ b/goblin/gremlin_python_driver/driver.py @@ -48,10 +48,11 @@ class Driver: class AsyncResponseIter: - def __init__(self, ws, loop, conn): + def __init__(self, ws, loop, conn, *, force_close=True): self._ws = ws self._loop = loop self._conn = conn + self._force_close = force_close self._closed = False async def __aiter__(self): @@ -62,6 +63,8 @@ class AsyncResponseIter: if msg: return msg else: + if self._force_close: + await self.close() raise StopAsyncIteration async def close(self): diff --git a/tests/test_engine.py b/tests/test_engine.py index 4081dd6bf1726c94265f6edce035f0f112f9e2bc..c9bf716b40906a93ff62b05dec57bc24b1ca095c 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -85,7 +85,6 @@ class TestEngine(unittest.TestCase): results += msg.data print(len(results)) self.assertTrue(len(results) > 1) - await stream.close() await engine.close() self.loop.run_until_complete(go())