Commit e18200f8 authored by davebshow's avatar davebshow
Browse files

fixed bug with context manager close

parent 329fbaf5
......@@ -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)
......
......@@ -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
......
......@@ -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):
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment