From fe64744d4eb7d1b989b24d6178929793de57a7d5 Mon Sep 17 00:00:00 2001
From: Cory Johns <johnsca@gmail.com>
Date: Tue, 7 Mar 2017 16:36:22 -0600
Subject: [PATCH] Improve error handling

---
 juju/client/connection.py     | 18 +++++++++++++-----
 juju/utils.py                 |  6 ++++++
 tests/unit/test_connection.py |  4 ++++
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/juju/client/connection.py b/juju/client/connection.py
index 82ec4fa..b9eb3bc 100644
--- a/juju/client/connection.py
+++ b/juju/client/connection.py
@@ -84,14 +84,22 @@ class Connection:
         await self.ws.close()
 
     async def recv(self, request_id):
+        if not self.is_open:
+            raise websockets.exceptions.ConnectionClosed(0, 'websocket closed')
         return await self.messages.get(request_id)
 
     async def receiver(self):
         while self.is_open:
-            result = await self.ws.recv()
-            if result is not None:
-                result = json.loads(result)
-                await self.messages.put(result['request-id'], result)
+            try:
+                result = await self.ws.recv()
+                if result is not None:
+                    result = json.loads(result)
+                    await self.messages.put(result['request-id'], result)
+            except Exception as e:
+                await self.messages.put_all(e)
+                raise
+        await self.messages.put_all(websockets.exceptions.ConnectionClosed(
+            0, 'websocket closed'))
 
     async def rpc(self, msg, encoder=None):
         self.__request_id__ += 1
@@ -111,7 +119,7 @@ class Connection:
             # API Error Response
             raise JujuAPIError(result)
 
-        if not 'response' in result:
+        if 'response' not in result:
             # This may never happen
             return result
 
diff --git a/juju/utils.py b/juju/utils.py
index 9f5d63d..f4db66e 100644
--- a/juju/utils.py
+++ b/juju/utils.py
@@ -59,7 +59,13 @@ class IdQueue:
     async def get(self, id):
         value = await self._queues[id].get()
         del self._queues[id]
+        if isinstance(value, Exception):
+            raise value
         return value
 
     async def put(self, id, value):
         await self._queues[id].put(value)
+
+    async def put_all(self, value):
+        for queue in self._queues.values():
+            await queue.put(value)
diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py
index 354b363..5371fdb 100644
--- a/tests/unit/test_connection.py
+++ b/tests/unit/test_connection.py
@@ -3,6 +3,8 @@ import mock
 import pytest
 from collections import deque
 
+from websockets.exceptions import ConnectionClosed
+
 from .. import base
 from juju.client.connection import Connection
 
@@ -17,6 +19,8 @@ class WebsocketMock:
         pass
 
     async def recv(self):
+        if not self.responses:
+            raise ConnectionClosed(0, 'no reason')
         return json.dumps(self.responses.popleft())
 
     async def close(self):
-- 
2.17.1