From afd9bf49afc5025f3b668fd3a09b0275d7edc382 Mon Sep 17 00:00:00 2001
From: Cory Johns <johnsca@gmail.com>
Date: Tue, 7 Mar 2017 16:04:37 -0600
Subject: [PATCH] Added test for out of order receive and fix bug creating task

---
 juju/client/connection.py     |  2 +-
 tests/base.py                 |  5 ++++
 tests/unit/__init__.py        |  0
 tests/unit/test_connection.py | 45 +++++++++++++++++++++++++++++++++++
 4 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 tests/unit/__init__.py
 create mode 100644 tests/unit/test_connection.py

diff --git a/juju/client/connection.py b/juju/client/connection.py
index 625c609..82ec4fa 100644
--- a/juju/client/connection.py
+++ b/juju/client/connection.py
@@ -76,6 +76,7 @@ class Connection:
         kw['loop'] = self.loop
         self.addr = url
         self.ws = await websockets.connect(url, **kw)
+        self.loop.create_task(self.receiver())
         log.info("Driver connected to juju %s", url)
         return self
 
@@ -216,7 +217,6 @@ class Connection:
         client = cls(endpoint, uuid, username, password, cacert, macaroons,
                      loop)
         await client.open()
-        client.loop.create_task(client.receiver)
 
         redirect_info = await client.redirect_info()
         if not redirect_info:
diff --git a/tests/base.py b/tests/base.py
index af386ea..292d04a 100644
--- a/tests/base.py
+++ b/tests/base.py
@@ -44,3 +44,8 @@ class CleanModel():
         await self.model.disconnect()
         await self.controller.destroy_model(self.model.info.uuid)
         await self.controller.disconnect()
+
+
+class AsyncMock(mock.MagicMock):
+    async def __call__(self, *args, **kwargs):
+        return super().__call__(*args, **kwargs)
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py
new file mode 100644
index 0000000..354b363
--- /dev/null
+++ b/tests/unit/test_connection.py
@@ -0,0 +1,45 @@
+import json
+import mock
+import pytest
+from collections import deque
+
+from .. import base
+from juju.client.connection import Connection
+
+
+class WebsocketMock:
+    def __init__(self, responses):
+        super().__init__()
+        self.responses = deque(responses)
+        self.open = True
+
+    async def send(self, message):
+        pass
+
+    async def recv(self):
+        return json.dumps(self.responses.popleft())
+
+    async def close(self):
+        self.open = False
+
+
+@pytest.mark.asyncio
+async def test_out_of_order(event_loop):
+    con = Connection(*[None]*4)
+    ws = WebsocketMock([
+        {'request-id': 1},
+        {'request-id': 3},
+        {'request-id': 2},
+    ])
+    expected_responses = [
+        {'request-id': 1},
+        {'request-id': 2},
+        {'request-id': 3},
+    ]
+    con._get_sll = mock.MagicMock()
+    with mock.patch('websockets.connect', base.AsyncMock(return_value=ws)):
+        await con.open()
+    actual_responses = []
+    for i in range(3):
+        actual_responses.append(await con.rpc({'version': 1}))
+    assert actual_responses == expected_responses
-- 
2.25.1