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