tests/unit/test_connection.py

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