From 3b4fc01827ffa9e78a1cb42e999f71156a58124c Mon Sep 17 00:00:00 2001
From: davebshow <davebshow@gmail.com>
Date: Wed, 6 Jul 2016 13:29:01 -0400
Subject: [PATCH] fixed a couple small problems with #1
---
goblin/driver/api.py | 1 -
goblin/driver/connection.py | 14 ++++++--------
2 files changed, 6 insertions(+), 9 deletions(-)
diff --git a/goblin/driver/api.py b/goblin/driver/api.py
index 599d49b..24154de 100644
--- a/goblin/driver/api.py
+++ b/goblin/driver/api.py
@@ -1,7 +1,6 @@
import asyncio
import aiohttp
-from goblin.driver import pool
from goblin.driver import connection
diff --git a/goblin/driver/connection.py b/goblin/driver/connection.py
index 82d4e3a..f4bd98a 100644
--- a/goblin/driver/connection.py
+++ b/goblin/driver/connection.py
@@ -34,11 +34,9 @@ class Response:
async def fetch_data(self):
if self._done:
return None
-
msg = await self._response_queue.get()
- if msg is None: # end of response sentinel
+ if msg is None:
self._done = True
-
return msg
@@ -115,7 +113,7 @@ class Connection(AbstractConnection):
self.response_queues[request_id] = response_queue
self._ws.send_bytes(message)
self._loop.create_task(self.receive())
- return Response(response_queue, self._loop, self)
+ return Response(response_queue, self._loop)
async def close(self):
await self._ws.close()
@@ -181,19 +179,19 @@ class Connection(AbstractConnection):
message["status"]["message"],
message["result"]["meta"])
response_queue = self._response_queues[request_id]
- if message.status_code not in (206, 407):
- # this message concludes the response
- await response_queue.put(None)
- del self._response_queues[request_id]
if message.status_code in [200, 206, 204]:
response_queue.put_nowait(message)
if message.status_code == 206:
self._loop.create_task(self.receive())
+ else:
+ response_queue.put_nowait(None)
+ del self._response_queues[request_id]
elif message.status_code == 407:
self._authenticate(self._username, self._password,
self._processor, self._session)
self._loop.create_task(self.receive())
else:
+ del self._response_queues[request_id]
raise RuntimeError("{0} {1}".format(message.status_code,
message.message))
--
GitLab