From e18200f8059cd1e05ddca0985c99eba0404f76dd Mon Sep 17 00:00:00 2001 From: davebshow <davebshow@gmail.com> Date: Sun, 17 May 2015 19:48:15 -0400 Subject: [PATCH] fixed bug with context manager close --- aiogremlin/connection.py | 6 +++--- aiogremlin/contextmanager.py | 1 + tests/tests.py | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/aiogremlin/connection.py b/aiogremlin/connection.py index a7c75dc..7f265d1 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 5e97e25..9e7e9a2 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 0c396d6..6896c02 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): -- GitLab