fixed bug with context manager close

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)

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

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):