X-Git-Url: https://osm.etsi.org/gitweb/?a=blobdiff_plain;f=tests%2Funit%2Ftest_connection.py;h=f69b8d6bc2a0a102957f9d1d0a29b59933110472;hb=dbcac620b62e16f9e4b7a763e1400f2f3a79b285;hp=354b363673b8d86385a678d0e1f61b9f9f38de07;hpb=afd9bf49afc5025f3b668fd3a09b0275d7edc382;p=osm%2FN2VC.git

diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py
index 354b363..f69b8d6 100644
--- a/tests/unit/test_connection.py
+++ b/tests/unit/test_connection.py
@@ -1,8 +1,11 @@
+import asyncio
 import json
 import mock
 import pytest
 from collections import deque
 
+from websockets.exceptions import ConnectionClosed
+
 from .. import base
 from juju.client.connection import Connection
 
@@ -17,6 +20,9 @@ class WebsocketMock:
         pass
 
     async def recv(self):
+        if not self.responses:
+            await asyncio.sleep(1)  # delay to give test time to finish
+            raise ConnectionClosed(0, 'ran out of responses')
         return json.dumps(self.responses.popleft())
 
     async def close(self):
@@ -37,9 +43,12 @@ async def test_out_of_order(event_loop):
         {'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
+    try:
+        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
+    finally:
+        await con.close()