diff --git a/aiogremlin/connection.py b/aiogremlin/connection.py index a7c75dc2cdaf06dd9c0cfdb6bb58ebebcf3c236a..7f265d1aa7a336e236a2a653a7e6d5df614d2555 100644 --- a/aiogremlin/connection.py +++ b/aiogremlin/connection.py @@ -134,8 +134,8 @@ class GremlinClientWebSocketResponse(BaseConnection, ClientWebSocketResponse): @asyncio.coroutine def close(self, *, code=1000, message=b''): if not self._closed: - closed = self._close() - if closed: + do_close = self._close() + if do_close: return True while True: try: @@ -158,7 +158,7 @@ class GremlinClientWebSocketResponse(BaseConnection, ClientWebSocketResponse): else: return False - def _close(self): + def _close(self, code=1000, message=b''): self._closed = True try: self._writer.close(code, message) diff --git a/aiogremlin/contextmanager.py b/aiogremlin/contextmanager.py index 5e97e2544810d8fd48f5543252945a9cc9c2fb7a..9e7e9a2f2168eb6f820c360ec5bc2df5e9fbed12 100644 --- a/aiogremlin/contextmanager.py +++ b/aiogremlin/contextmanager.py @@ -11,6 +11,7 @@ class ConnectionContextManager: def __exit__(self, exception_type, exception_value, traceback): try: + self._conn._closing = True self._conn._close() finally: self._conn = None diff --git a/tests/tests.py b/tests/tests.py index 0c396d6706b162c63a0b7d2e03e1c0fad40f61fb..6896c0296c48bd75f7b9ae3a90717622d4695e27 100644 --- a/tests/tests.py +++ b/tests/tests.py @@ -229,7 +229,7 @@ class ContextMngrTest(unittest.TestCase): factory=AiohttpFactory, max_retries=0) def tearDown(self): - # self.loop.run_until_complete(self.pool.close()) + self.loop.run_until_complete(self.pool.close()) self.loop.close() def test_connection_manager(self):