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