modules/libjuju/tests/unit/test_connection.py

   1 import asyncio
   2 import json
   3 import mock
   4 import pytest
   5 from collections import deque
   6
   7 from websockets.exceptions import ConnectionClosed
   8
   9 from .. import base
  10 from juju.client.connection import Connection
  11
  12
  13 class WebsocketMock:
  14     def __init__(self, responses):
  15         super().__init__()
  16         self.responses = deque(responses)
  17         self.open = True
  18
  19     async def send(self, message):
  20         pass
  21
  22     async def recv(self):
  23         if not self.responses:
  24             await asyncio.sleep(1)  # delay to give test time to finish
  25             raise ConnectionClosed(0, 'ran out of responses')
  26         return json.dumps(self.responses.popleft())
  27
  28     async def close(self):
  29         self.open = False
  30
  31
  32 @pytest.mark.asyncio
  33 async def test_out_of_order(event_loop):
  34     con = Connection(*[None]*4)
  35     ws = WebsocketMock([
  36         {'request-id': 1},
  37         {'request-id': 3},
  38         {'request-id': 2},
  39     ])
  40     expected_responses = [
  41         {'request-id': 1},
  42         {'request-id': 2},
  43         {'request-id': 3},
  44     ]
  45     con._get_sll = mock.MagicMock()
  46     try:
  47         with mock.patch('websockets.connect', base.AsyncMock(return_value=ws)):
  48             await con.open()
  49         actual_responses = []
  50         for i in range(3):
  51             actual_responses.append(await con.rpc({'version': 1}))
  52         assert actual_responses == expected_responses
  53     finally:
  54         await con.close()