From 2399bdd31415f9be498cd0e6f805586d7c18015d Mon Sep 17 00:00:00 2001
From: Pete Vander Giessen <petevg@gmail.com>
Date: Tue, 7 Mar 2017 14:54:49 -0600
Subject: [PATCH] Fix issue where we do not check to make sure that we are
 receiving the correct response.

---
 juju/client/connection.py | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/juju/client/connection.py b/juju/client/connection.py
index b508a1a..486a57f 100644
--- a/juju/client/connection.py
+++ b/juju/client/connection.py
@@ -52,6 +52,7 @@ class Connection:
         self.addr = None
         self.ws = None
         self.facades = {}
+        self.messages = {}
 
     @property
     def is_open(self):
@@ -80,12 +81,22 @@ class Connection:
     async def close(self):
         await self.ws.close()
 
-    async def recv(self):
-        result = await self.ws.recv()
-        if result is not None:
-            result = json.loads(result)
+    async def recv(self, request_id):
+        while not self.messages.get(request_id):
+            await asyncio.sleep(0)
+
+        result = self.messages[request_id]
+
+        del self.messages[request_id]
         return result
 
+    async def receiver(self):
+        while self.is_open:
+            result = await self.ws.recv()
+            if result is not None:
+                result = json.loads(result)
+                self.messages[result['request-id']] = result
+
     async def rpc(self, msg, encoder=None):
         self.__request_id__ += 1
         msg['request-id'] = self.__request_id__
@@ -95,7 +106,7 @@ class Connection:
             msg['version'] = self.facades[msg['type']]
         outgoing = json.dumps(msg, indent=2, cls=encoder)
         await self.ws.send(outgoing)
-        result = await self.recv()
+        result = await self.recv(msg['request-id'])
 
         if not result:
             return result
@@ -210,6 +221,7 @@ class Connection:
         client = cls(endpoint, uuid, username, password, cacert, macaroons,
                      loop)
         await client.open()
+        self.loop.create_task(self.receiver)
 
         redirect_info = await client.redirect_info()
         if not redirect_info:
-- 
2.17.1