Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

1216 строки
45KB

  1. import asyncio
  2. import sys
  3. import time
  4. import unittest
  5. import six
  6. if six.PY3:
  7. from unittest import mock
  8. else:
  9. import mock
  10. try:
  11. import websockets
  12. except ImportError:
  13. # weirdness to avoid errors in PY2 test run
  14. class _dummy():
  15. pass
  16. websockets = _dummy()
  17. websockets.exceptions = _dummy()
  18. websockets.exceptions.InvalidURI = _dummy()
  19. from engineio import asyncio_client
  20. from engineio import client
  21. from engineio import exceptions
  22. from engineio import packet
  23. from engineio import payload
  24. def AsyncMock(*args, **kwargs):
  25. """Return a mock asynchronous function."""
  26. m = mock.MagicMock(*args, **kwargs)
  27. async def mock_coro(*args, **kwargs):
  28. return m(*args, **kwargs)
  29. mock_coro.mock = m
  30. return mock_coro
  31. def _run(coro):
  32. """Run the given coroutine."""
  33. return asyncio.get_event_loop().run_until_complete(coro)
  34. @unittest.skipIf(sys.version_info < (3, 5), 'only for Python 3.5+')
  35. class TestAsyncClient(unittest.TestCase):
  36. def test_is_asyncio_based(self):
  37. c = asyncio_client.AsyncClient()
  38. self.assertEqual(c.is_asyncio_based(), True)
  39. def test_already_connected(self):
  40. c = asyncio_client.AsyncClient()
  41. c.state = 'connected'
  42. self.assertRaises(ValueError, _run, c.connect('http://foo'))
  43. def test_invalid_transports(self):
  44. c = asyncio_client.AsyncClient()
  45. self.assertRaises(ValueError, _run, c.connect(
  46. 'http://foo', transports=['foo', 'bar']))
  47. def test_some_invalid_transports(self):
  48. c = asyncio_client.AsyncClient()
  49. c._connect_websocket = AsyncMock()
  50. _run(c.connect('http://foo', transports=['foo', 'websocket', 'bar']))
  51. self.assertEqual(c.transports, ['websocket'])
  52. def test_connect_polling(self):
  53. c = asyncio_client.AsyncClient()
  54. c._connect_polling = AsyncMock(return_value='foo')
  55. self.assertEqual(_run(c.connect('http://foo')), 'foo')
  56. c._connect_polling.mock.assert_called_once_with(
  57. 'http://foo', {}, 'engine.io')
  58. c = asyncio_client.AsyncClient()
  59. c._connect_polling = AsyncMock(return_value='foo')
  60. self.assertEqual(
  61. _run(c.connect('http://foo', transports=['polling'])), 'foo')
  62. c._connect_polling.mock.assert_called_once_with(
  63. 'http://foo', {}, 'engine.io')
  64. c = asyncio_client.AsyncClient()
  65. c._connect_polling = AsyncMock(return_value='foo')
  66. self.assertEqual(
  67. _run(c.connect('http://foo', transports=['polling',
  68. 'websocket'])),
  69. 'foo')
  70. c._connect_polling.mock.assert_called_once_with(
  71. 'http://foo', {}, 'engine.io')
  72. def test_connect_websocket(self):
  73. c = asyncio_client.AsyncClient()
  74. c._connect_websocket = AsyncMock(return_value='foo')
  75. self.assertEqual(
  76. _run(c.connect('http://foo', transports=['websocket'])),
  77. 'foo')
  78. c._connect_websocket.mock.assert_called_once_with(
  79. 'http://foo', {}, 'engine.io')
  80. c = asyncio_client.AsyncClient()
  81. c._connect_websocket = AsyncMock(return_value='foo')
  82. self.assertEqual(
  83. _run(c.connect('http://foo', transports='websocket')),
  84. 'foo')
  85. c._connect_websocket.mock.assert_called_once_with(
  86. 'http://foo', {}, 'engine.io')
  87. def test_connect_query_string(self):
  88. c = asyncio_client.AsyncClient()
  89. c._connect_polling = AsyncMock(return_value='foo')
  90. self.assertEqual(_run(c.connect('http://foo?bar=baz')), 'foo')
  91. c._connect_polling.mock.assert_called_once_with(
  92. 'http://foo?bar=baz', {}, 'engine.io')
  93. def test_connect_custom_headers(self):
  94. c = asyncio_client.AsyncClient()
  95. c._connect_polling = AsyncMock(return_value='foo')
  96. self.assertEqual(
  97. _run(c.connect('http://foo', headers={'Foo': 'Bar'})),
  98. 'foo')
  99. c._connect_polling.mock.assert_called_once_with(
  100. 'http://foo', {'Foo': 'Bar'}, 'engine.io')
  101. def test_wait(self):
  102. c = asyncio_client.AsyncClient()
  103. done = []
  104. async def fake_read_look_task():
  105. done.append(True)
  106. c.read_loop_task = fake_read_look_task()
  107. _run(c.wait())
  108. self.assertEqual(done, [True])
  109. def test_wait_no_task(self):
  110. c = asyncio_client.AsyncClient()
  111. c.read_loop_task = None
  112. _run(c.wait())
  113. def test_send(self):
  114. c = asyncio_client.AsyncClient()
  115. saved_packets = []
  116. async def fake_send_packet(pkt):
  117. saved_packets.append(pkt)
  118. c._send_packet = fake_send_packet
  119. _run(c.send('foo'))
  120. _run(c.send('foo', binary=False))
  121. _run(c.send(b'foo', binary=True))
  122. self.assertEqual(saved_packets[0].packet_type, packet.MESSAGE)
  123. self.assertEqual(saved_packets[0].data, 'foo')
  124. self.assertEqual(saved_packets[0].binary,
  125. False if six.PY3 else True)
  126. self.assertEqual(saved_packets[1].packet_type, packet.MESSAGE)
  127. self.assertEqual(saved_packets[1].data, 'foo')
  128. self.assertEqual(saved_packets[1].binary, False)
  129. self.assertEqual(saved_packets[2].packet_type, packet.MESSAGE)
  130. self.assertEqual(saved_packets[2].data, b'foo')
  131. self.assertEqual(saved_packets[2].binary, True)
  132. def test_disconnect_not_connected(self):
  133. c = asyncio_client.AsyncClient()
  134. c.state = 'foo'
  135. c.sid = 'bar'
  136. _run(c.disconnect())
  137. self.assertEqual(c.state, 'disconnected')
  138. self.assertIsNone(c.sid)
  139. def test_disconnect_polling(self):
  140. c = asyncio_client.AsyncClient()
  141. client.connected_clients.append(c)
  142. c.state = 'connected'
  143. c.current_transport = 'polling'
  144. c.queue = mock.MagicMock()
  145. c.queue.put = AsyncMock()
  146. c.queue.join = AsyncMock()
  147. c.read_loop_task = AsyncMock()()
  148. c.ws = mock.MagicMock()
  149. c.ws.close = AsyncMock()
  150. c._trigger_event = AsyncMock()
  151. _run(c.disconnect())
  152. c.ws.close.mock.assert_not_called()
  153. self.assertNotIn(c, client.connected_clients)
  154. c._trigger_event.mock.assert_called_once_with('disconnect',
  155. run_async=False)
  156. def test_disconnect_websocket(self):
  157. c = asyncio_client.AsyncClient()
  158. client.connected_clients.append(c)
  159. c.state = 'connected'
  160. c.current_transport = 'websocket'
  161. c.queue = mock.MagicMock()
  162. c.queue.put = AsyncMock()
  163. c.queue.join = AsyncMock()
  164. c.read_loop_task = AsyncMock()()
  165. c.ws = mock.MagicMock()
  166. c.ws.close = AsyncMock()
  167. c._trigger_event = AsyncMock()
  168. _run(c.disconnect())
  169. c.ws.close.mock.assert_called_once_with()
  170. self.assertNotIn(c, client.connected_clients)
  171. c._trigger_event.mock.assert_called_once_with('disconnect',
  172. run_async=False)
  173. def test_disconnect_polling_abort(self):
  174. c = asyncio_client.AsyncClient()
  175. client.connected_clients.append(c)
  176. c.state = 'connected'
  177. c.current_transport = 'polling'
  178. c.queue = mock.MagicMock()
  179. c.queue.put = AsyncMock()
  180. c.queue.join = AsyncMock()
  181. c.read_loop_task = AsyncMock()()
  182. c.ws = mock.MagicMock()
  183. c.ws.close = AsyncMock()
  184. _run(c.disconnect(abort=True))
  185. c.queue.join.mock.assert_not_called()
  186. c.ws.close.mock.assert_not_called()
  187. self.assertNotIn(c, client.connected_clients)
  188. def test_disconnect_websocket_abort(self):
  189. c = asyncio_client.AsyncClient()
  190. client.connected_clients.append(c)
  191. c.state = 'connected'
  192. c.current_transport = 'websocket'
  193. c.queue = mock.MagicMock()
  194. c.queue.put = AsyncMock()
  195. c.queue.join = AsyncMock()
  196. c.read_loop_task = AsyncMock()()
  197. c.ws = mock.MagicMock()
  198. c.ws.close = AsyncMock()
  199. _run(c.disconnect(abort=True))
  200. c.queue.join.mock.assert_not_called()
  201. c.ws.mock.assert_not_called()
  202. self.assertNotIn(c, client.connected_clients)
  203. def test_background_tasks(self):
  204. r = []
  205. async def foo(arg):
  206. r.append(arg)
  207. c = asyncio_client.AsyncClient()
  208. c.start_background_task(foo, 'bar')
  209. pending = asyncio.Task.all_tasks()
  210. asyncio.get_event_loop().run_until_complete(asyncio.wait(pending))
  211. self.assertEqual(r, ['bar'])
  212. def test_sleep(self):
  213. c = asyncio_client.AsyncClient()
  214. _run(c.sleep(0))
  215. def test_create_queue(self):
  216. c = asyncio_client.AsyncClient()
  217. q = c.create_queue()
  218. self.assertRaises(q.Empty, q.get_nowait)
  219. def test_create_event(self):
  220. c = asyncio_client.AsyncClient()
  221. e = c.create_event()
  222. self.assertFalse(e.is_set())
  223. e.set()
  224. self.assertTrue(e.is_set())
  225. @mock.patch('engineio.client.time.time', return_value=123.456)
  226. def test_polling_connection_failed(self, _time):
  227. c = asyncio_client.AsyncClient()
  228. c._send_request = AsyncMock(return_value=None)
  229. self.assertRaises(
  230. exceptions.ConnectionError, _run, c.connect(
  231. 'http://foo', headers={'Foo': 'Bar'}))
  232. c._send_request.mock.assert_called_once_with(
  233. 'GET', 'http://foo/engine.io/?transport=polling&EIO=3&t=123.456',
  234. headers={'Foo': 'Bar'})
  235. def test_polling_connection_404(self):
  236. c = asyncio_client.AsyncClient()
  237. c._send_request = AsyncMock()
  238. c._send_request.mock.return_value.status = 404
  239. self.assertRaises(
  240. exceptions.ConnectionError, _run, c.connect('http://foo'))
  241. def test_polling_connection_invalid_packet(self):
  242. c = asyncio_client.AsyncClient()
  243. c._send_request = AsyncMock()
  244. c._send_request.mock.return_value.status = 200
  245. c._send_request.mock.return_value.read = AsyncMock(
  246. return_value=b'foo')
  247. self.assertRaises(
  248. exceptions.ConnectionError, _run, c.connect('http://foo'))
  249. def test_polling_connection_no_open_packet(self):
  250. c = asyncio_client.AsyncClient()
  251. c._send_request = AsyncMock()
  252. c._send_request.mock.return_value.status = 200
  253. c._send_request.mock.return_value.read = AsyncMock(
  254. return_value=payload.Payload(packets=[
  255. packet.Packet(packet.CLOSE, {
  256. 'sid': '123', 'upgrades': [], 'pingInterval': 10,
  257. 'pingTimeout': 20
  258. })
  259. ]).encode())
  260. self.assertRaises(
  261. exceptions.ConnectionError, _run, c.connect('http://foo'))
  262. def test_polling_connection_successful(self):
  263. c = asyncio_client.AsyncClient()
  264. c._send_request = AsyncMock()
  265. c._send_request.mock.return_value.status = 200
  266. c._send_request.mock.return_value.read = AsyncMock(
  267. return_value=payload.Payload(packets=[
  268. packet.Packet(packet.OPEN, {
  269. 'sid': '123', 'upgrades': [], 'pingInterval': 1000,
  270. 'pingTimeout': 2000
  271. })
  272. ]).encode())
  273. c._ping_loop = AsyncMock()
  274. c._read_loop_polling = AsyncMock()
  275. c._read_loop_websocket = AsyncMock()
  276. c._write_loop = AsyncMock()
  277. on_connect = AsyncMock()
  278. c.on('connect', on_connect)
  279. _run(c.connect('http://foo'))
  280. time.sleep(0.1)
  281. c._ping_loop.mock.assert_called_once_with()
  282. c._read_loop_polling.mock.assert_called_once_with()
  283. c._read_loop_websocket.mock.assert_not_called()
  284. c._write_loop.mock.assert_called_once_with()
  285. on_connect.mock.assert_called_once_with()
  286. self.assertIn(c, client.connected_clients)
  287. self.assertEqual(
  288. c.base_url,
  289. 'http://foo/engine.io/?transport=polling&EIO=3&sid=123')
  290. self.assertEqual(c.sid, '123')
  291. self.assertEqual(c.ping_interval, 1)
  292. self.assertEqual(c.ping_timeout, 2)
  293. self.assertEqual(c.upgrades, [])
  294. self.assertEqual(c.transport(), 'polling')
  295. def test_polling_connection_with_more_packets(self):
  296. c = asyncio_client.AsyncClient()
  297. c._send_request = AsyncMock()
  298. c._send_request.mock.return_value.status = 200
  299. c._send_request.mock.return_value.read = AsyncMock(
  300. return_value=payload.Payload(packets=[
  301. packet.Packet(packet.OPEN, {
  302. 'sid': '123', 'upgrades': [], 'pingInterval': 1000,
  303. 'pingTimeout': 2000
  304. }),
  305. packet.Packet(packet.NOOP)
  306. ]).encode())
  307. c._ping_loop = AsyncMock()
  308. c._read_loop_polling = AsyncMock()
  309. c._read_loop_websocket = AsyncMock()
  310. c._write_loop = AsyncMock()
  311. c._receive_packet = AsyncMock()
  312. on_connect = AsyncMock()
  313. c.on('connect', on_connect)
  314. _run(c.connect('http://foo'))
  315. time.sleep(0.1)
  316. self.assertEqual(c._receive_packet.mock.call_count, 1)
  317. self.assertEqual(
  318. c._receive_packet.mock.call_args_list[0][0][0].packet_type,
  319. packet.NOOP)
  320. def test_polling_connection_upgraded(self):
  321. c = asyncio_client.AsyncClient()
  322. c._send_request = AsyncMock()
  323. c._send_request.mock.return_value.status = 200
  324. c._send_request.mock.return_value.read = AsyncMock(
  325. return_value=payload.Payload(packets=[
  326. packet.Packet(packet.OPEN, {
  327. 'sid': '123', 'upgrades': ['websocket'],
  328. 'pingInterval': 1000, 'pingTimeout': 2000
  329. })
  330. ]).encode())
  331. c._connect_websocket = AsyncMock(return_value=True)
  332. on_connect = mock.MagicMock()
  333. c.on('connect', on_connect)
  334. _run(c.connect('http://foo'))
  335. c._connect_websocket.mock.assert_called_once_with('http://foo', {},
  336. 'engine.io')
  337. on_connect.assert_called_once_with()
  338. self.assertIn(c, client.connected_clients)
  339. self.assertEqual(
  340. c.base_url,
  341. 'http://foo/engine.io/?transport=polling&EIO=3&sid=123')
  342. self.assertEqual(c.sid, '123')
  343. self.assertEqual(c.ping_interval, 1)
  344. self.assertEqual(c.ping_timeout, 2)
  345. self.assertEqual(c.upgrades, ['websocket'])
  346. def test_polling_connection_not_upgraded(self):
  347. c = asyncio_client.AsyncClient()
  348. c._send_request = AsyncMock()
  349. c._send_request.mock.return_value.status = 200
  350. c._send_request.mock.return_value.read = AsyncMock(
  351. return_value=payload.Payload(packets=[
  352. packet.Packet(packet.OPEN, {
  353. 'sid': '123', 'upgrades': ['websocket'],
  354. 'pingInterval': 1000, 'pingTimeout': 2000
  355. })
  356. ]).encode())
  357. c._connect_websocket = AsyncMock(return_value=False)
  358. c._ping_loop = AsyncMock()
  359. c._read_loop_polling = AsyncMock()
  360. c._read_loop_websocket = AsyncMock()
  361. c._write_loop = AsyncMock()
  362. on_connect = mock.MagicMock()
  363. c.on('connect', on_connect)
  364. _run(c.connect('http://foo'))
  365. time.sleep(0.1)
  366. c._connect_websocket.mock.assert_called_once_with('http://foo', {},
  367. 'engine.io')
  368. c._ping_loop.mock.assert_called_once_with()
  369. c._read_loop_polling.mock.assert_called_once_with()
  370. c._read_loop_websocket.mock.assert_not_called()
  371. c._write_loop.mock.assert_called_once_with()
  372. on_connect.assert_called_once_with()
  373. self.assertIn(c, client.connected_clients)
  374. @mock.patch('engineio.client.time.time', return_value=123.456)
  375. @mock.patch('engineio.asyncio_client.websockets.connect', new=AsyncMock(
  376. side_effect=[websockets.exceptions.InvalidURI]))
  377. def test_websocket_connection_failed(self, _time):
  378. c = asyncio_client.AsyncClient()
  379. self.assertRaises(
  380. exceptions.ConnectionError, _run,
  381. c.connect('http://foo', transports=['websocket'],
  382. headers={'Foo': 'Bar'}))
  383. asyncio_client.websockets.connect.mock.assert_called_once_with(
  384. 'ws://foo/engine.io/?transport=websocket&EIO=3&t=123.456',
  385. extra_headers={'Foo': 'Bar'})
  386. @mock.patch('engineio.client.time.time', return_value=123.456)
  387. @mock.patch('engineio.asyncio_client.websockets.connect', new=AsyncMock(
  388. side_effect=[websockets.exceptions.InvalidURI]))
  389. def test_websocket_upgrade_failed(self, _time):
  390. c = asyncio_client.AsyncClient()
  391. c.sid = '123'
  392. self.assertFalse(_run(c.connect(
  393. 'http://foo', transports=['websocket'])))
  394. asyncio_client.websockets.connect.mock.assert_called_once_with(
  395. 'ws://foo/engine.io/?transport=websocket&EIO=3&sid=123&t=123.456',
  396. extra_headers={})
  397. @mock.patch('engineio.asyncio_client.websockets.connect', new=AsyncMock())
  398. def test_websocket_connection_no_open_packet(self):
  399. asyncio_client.websockets.connect.mock.return_value.recv = AsyncMock(
  400. return_value=packet.Packet(packet.CLOSE).encode())
  401. c = asyncio_client.AsyncClient()
  402. self.assertRaises(
  403. exceptions.ConnectionError, _run,
  404. c.connect('http://foo', transports=['websocket']))
  405. @mock.patch('engineio.client.time.time', return_value=123.456)
  406. @mock.patch('engineio.asyncio_client.websockets.connect', new=AsyncMock())
  407. def test_websocket_connection_successful(self, _time):
  408. ws = asyncio_client.websockets.connect.mock.return_value
  409. ws.recv = AsyncMock(return_value=packet.Packet(
  410. packet.OPEN, {
  411. 'sid': '123', 'upgrades': [], 'pingInterval': 1000,
  412. 'pingTimeout': 2000
  413. }).encode())
  414. c = asyncio_client.AsyncClient()
  415. c._ping_loop = AsyncMock()
  416. c._read_loop_polling = AsyncMock()
  417. c._read_loop_websocket = AsyncMock()
  418. c._write_loop = AsyncMock()
  419. on_connect = mock.MagicMock()
  420. c.on('connect', on_connect)
  421. _run(c.connect('ws://foo', transports=['websocket']))
  422. time.sleep(0.1)
  423. c._ping_loop.mock.assert_called_once_with()
  424. c._read_loop_polling.mock.assert_not_called()
  425. c._read_loop_websocket.mock.assert_called_once_with()
  426. c._write_loop.mock.assert_called_once_with()
  427. on_connect.assert_called_once_with()
  428. self.assertIn(c, client.connected_clients)
  429. self.assertEqual(
  430. c.base_url,
  431. 'ws://foo/engine.io/?transport=websocket&EIO=3')
  432. self.assertEqual(c.sid, '123')
  433. self.assertEqual(c.ping_interval, 1)
  434. self.assertEqual(c.ping_timeout, 2)
  435. self.assertEqual(c.upgrades, [])
  436. self.assertEqual(c.transport(), 'websocket')
  437. self.assertEqual(c.ws, ws)
  438. asyncio_client.websockets.connect.mock.assert_called_once_with(
  439. 'ws://foo/engine.io/?transport=websocket&EIO=3&t=123.456',
  440. extra_headers={})
  441. @mock.patch('engineio.client.time.time', return_value=123.456)
  442. @mock.patch('engineio.asyncio_client.websockets.connect', new=AsyncMock())
  443. def test_websocket_connection_with_cookies(self, _time):
  444. ws = asyncio_client.websockets.connect.mock.return_value
  445. ws.recv = AsyncMock(return_value=packet.Packet(
  446. packet.OPEN, {
  447. 'sid': '123', 'upgrades': [], 'pingInterval': 1000,
  448. 'pingTimeout': 2000
  449. }).encode())
  450. c = asyncio_client.AsyncClient()
  451. c.http = mock.MagicMock()
  452. c.http._cookie_jar = [mock.MagicMock(), mock.MagicMock()]
  453. c.http._cookie_jar[0].key = 'key'
  454. c.http._cookie_jar[0].value = 'value'
  455. c.http._cookie_jar[1].key = 'key2'
  456. c.http._cookie_jar[1].value = 'value2'
  457. c._ping_loop = AsyncMock()
  458. c._read_loop_polling = AsyncMock()
  459. c._read_loop_websocket = AsyncMock()
  460. c._write_loop = AsyncMock()
  461. on_connect = mock.MagicMock()
  462. c.on('connect', on_connect)
  463. _run(c.connect('ws://foo', transports=['websocket']))
  464. time.sleep(0.1)
  465. asyncio_client.websockets.connect.mock.assert_called_once_with(
  466. 'ws://foo/engine.io/?transport=websocket&EIO=3&t=123.456',
  467. extra_headers={'Cookie': 'key=value; key2=value2'})
  468. @mock.patch('engineio.asyncio_client.websockets.connect', new=AsyncMock())
  469. def test_websocket_upgrade_no_pong(self):
  470. ws = asyncio_client.websockets.connect.mock.return_value
  471. ws.recv = AsyncMock(return_value=packet.Packet(
  472. packet.OPEN, {
  473. 'sid': '123', 'upgrades': [], 'pingInterval': 1000,
  474. 'pingTimeout': 2000
  475. }).encode())
  476. ws.send = AsyncMock()
  477. c = asyncio_client.AsyncClient()
  478. c.sid = '123'
  479. c.current_transport = 'polling'
  480. c._ping_loop = AsyncMock()
  481. c._read_loop_polling = AsyncMock()
  482. c._read_loop_websocket = AsyncMock()
  483. c._write_loop = AsyncMock()
  484. on_connect = mock.MagicMock()
  485. c.on('connect', on_connect)
  486. self.assertFalse(_run(c.connect('ws://foo',
  487. transports=['websocket'])))
  488. c._ping_loop.mock.assert_not_called()
  489. c._read_loop_polling.mock.assert_not_called()
  490. c._read_loop_websocket.mock.assert_not_called()
  491. c._write_loop.mock.assert_not_called()
  492. on_connect.assert_not_called()
  493. self.assertEqual(c.transport(), 'polling')
  494. ws.send.mock.assert_called_once_with('2probe')
  495. @mock.patch('engineio.asyncio_client.websockets.connect', new=AsyncMock())
  496. def test_websocket_upgrade_successful(self):
  497. ws = asyncio_client.websockets.connect.mock.return_value
  498. ws.recv = AsyncMock(return_value=packet.Packet(
  499. packet.PONG, 'probe').encode())
  500. ws.send = AsyncMock()
  501. c = asyncio_client.AsyncClient()
  502. c.sid = '123'
  503. c.base_url = 'http://foo'
  504. c.current_transport = 'polling'
  505. c._ping_loop = AsyncMock()
  506. c._read_loop_polling = AsyncMock()
  507. c._read_loop_websocket = AsyncMock()
  508. c._write_loop = AsyncMock()
  509. on_connect = mock.MagicMock()
  510. c.on('connect', on_connect)
  511. self.assertTrue(_run(c.connect('ws://foo',
  512. transports=['websocket'])))
  513. time.sleep(0.1)
  514. c._ping_loop.mock.assert_called_once_with()
  515. c._read_loop_polling.mock.assert_not_called()
  516. c._read_loop_websocket.mock.assert_called_once_with()
  517. c._write_loop.mock.assert_called_once_with()
  518. on_connect.assert_not_called() # was called by polling
  519. self.assertNotIn(c, client.connected_clients) # was added by polling
  520. self.assertEqual(c.base_url, 'http://foo') # not changed
  521. self.assertEqual(c.sid, '123') # not changed
  522. self.assertEqual(c.transport(), 'websocket')
  523. self.assertEqual(c.ws, ws)
  524. self.assertEqual(
  525. ws.send.mock.call_args_list[0],
  526. (('2probe',),)) # ping
  527. self.assertEqual(
  528. ws.send.mock.call_args_list[1],
  529. (('5',),)) # upgrade
  530. def test_receive_unknown_packet(self):
  531. c = asyncio_client.AsyncClient()
  532. _run(c._receive_packet(packet.Packet(encoded_packet=b'9')))
  533. # should be ignored
  534. def test_receive_noop_packet(self):
  535. c = asyncio_client.AsyncClient()
  536. _run(c._receive_packet(packet.Packet(packet.NOOP)))
  537. # should be ignored
  538. def test_receive_pong_packet(self):
  539. c = asyncio_client.AsyncClient()
  540. c.pong_received = False
  541. _run(c._receive_packet(packet.Packet(packet.PONG)))
  542. self.assertTrue(c.pong_received)
  543. def test_receive_message_packet(self):
  544. c = asyncio_client.AsyncClient()
  545. c._trigger_event = AsyncMock()
  546. _run(c._receive_packet(packet.Packet(packet.MESSAGE, {'foo': 'bar'})))
  547. c._trigger_event.mock.assert_called_once_with(
  548. 'message', {'foo': 'bar'}, run_async=True)
  549. def test_receive_close_packet(self):
  550. c = asyncio_client.AsyncClient()
  551. c.disconnect = AsyncMock()
  552. _run(c._receive_packet(packet.Packet(packet.CLOSE)))
  553. c.disconnect.mock.assert_called_once_with(abort=True)
  554. def test_send_packet_disconnected(self):
  555. c = asyncio_client.AsyncClient()
  556. c.queue = c.create_queue()
  557. c.state = 'disconnected'
  558. _run(c._send_packet(packet.Packet(packet.NOOP)))
  559. self.assertTrue(c.queue.empty())
  560. def test_send_packet(self):
  561. c = asyncio_client.AsyncClient()
  562. c.queue = c.create_queue()
  563. c.state = 'connected'
  564. _run(c._send_packet(packet.Packet(packet.NOOP)))
  565. self.assertFalse(c.queue.empty())
  566. pkt = _run(c.queue.get())
  567. self.assertEqual(pkt.packet_type, packet.NOOP)
  568. def test_trigger_event_function(self):
  569. result = []
  570. def foo_handler(arg):
  571. result.append('ok')
  572. result.append(arg)
  573. c = asyncio_client.AsyncClient()
  574. c.on('message', handler=foo_handler)
  575. _run(c._trigger_event('message', 'bar'))
  576. self.assertEqual(result, ['ok', 'bar'])
  577. def test_trigger_event_coroutine(self):
  578. result = []
  579. async def foo_handler(arg):
  580. result.append('ok')
  581. result.append(arg)
  582. c = asyncio_client.AsyncClient()
  583. c.on('message', handler=foo_handler)
  584. _run(c._trigger_event('message', 'bar'))
  585. self.assertEqual(result, ['ok', 'bar'])
  586. def test_trigger_event_function_error(self):
  587. def connect_handler(arg):
  588. return 1 / 0
  589. def foo_handler(arg):
  590. return 1 / 0
  591. c = asyncio_client.AsyncClient()
  592. c.on('connect', handler=connect_handler)
  593. c.on('message', handler=foo_handler)
  594. self.assertFalse(_run(c._trigger_event('connect', '123')))
  595. self.assertIsNone(_run(c._trigger_event('message', 'bar')))
  596. def test_trigger_event_coroutine_error(self):
  597. async def connect_handler(arg):
  598. return 1 / 0
  599. async def foo_handler(arg):
  600. return 1 / 0
  601. c = asyncio_client.AsyncClient()
  602. c.on('connect', handler=connect_handler)
  603. c.on('message', handler=foo_handler)
  604. self.assertFalse(_run(c._trigger_event('connect', '123')))
  605. self.assertIsNone(_run(c._trigger_event('message', 'bar')))
  606. def test_trigger_event_function_async(self):
  607. result = []
  608. def foo_handler(arg):
  609. result.append('ok')
  610. result.append(arg)
  611. c = asyncio_client.AsyncClient()
  612. c.on('message', handler=foo_handler)
  613. fut = _run(c._trigger_event('message', 'bar', run_async=True))
  614. asyncio.get_event_loop().run_until_complete(fut)
  615. self.assertEqual(result, ['ok', 'bar'])
  616. def test_trigger_event_coroutine_async(self):
  617. result = []
  618. async def foo_handler(arg):
  619. result.append('ok')
  620. result.append(arg)
  621. c = asyncio_client.AsyncClient()
  622. c.on('message', handler=foo_handler)
  623. fut = _run(c._trigger_event('message', 'bar', run_async=True))
  624. asyncio.get_event_loop().run_until_complete(fut)
  625. self.assertEqual(result, ['ok', 'bar'])
  626. def test_trigger_event_function_async_error(self):
  627. result = []
  628. def foo_handler(arg):
  629. result.append(arg)
  630. return 1 / 0
  631. c = asyncio_client.AsyncClient()
  632. c.on('message', handler=foo_handler)
  633. fut = _run(c._trigger_event('message', 'bar', run_async=True))
  634. self.assertRaises(
  635. ZeroDivisionError, asyncio.get_event_loop().run_until_complete,
  636. fut)
  637. self.assertEqual(result, ['bar'])
  638. def test_trigger_event_coroutine_async_error(self):
  639. result = []
  640. async def foo_handler(arg):
  641. result.append(arg)
  642. return 1 / 0
  643. c = asyncio_client.AsyncClient()
  644. c.on('message', handler=foo_handler)
  645. fut = _run(c._trigger_event('message', 'bar', run_async=True))
  646. self.assertRaises(
  647. ZeroDivisionError, asyncio.get_event_loop().run_until_complete,
  648. fut)
  649. self.assertEqual(result, ['bar'])
  650. def test_trigger_unknown_event(self):
  651. c = asyncio_client.AsyncClient()
  652. _run(c._trigger_event('connect', run_async=False))
  653. _run(c._trigger_event('message', 123, run_async=True))
  654. # should do nothing
  655. def test_ping_loop_disconnected(self):
  656. c = asyncio_client.AsyncClient()
  657. c.state = 'disconnected'
  658. _run(c._ping_loop())
  659. # should not block
  660. def test_ping_loop_disconnect(self):
  661. c = asyncio_client.AsyncClient()
  662. c.state = 'connected'
  663. c.ping_interval = 10
  664. c._send_packet = AsyncMock()
  665. states = [
  666. ('disconnecting', True)
  667. ]
  668. async def fake_wait():
  669. c.state, c.pong_received = states.pop(0)
  670. c.ping_loop_event.wait = fake_wait
  671. _run(c._ping_loop())
  672. self.assertEqual(
  673. c._send_packet.mock.call_args_list[0][0][0].encode(), b'2')
  674. def test_ping_loop_missing_pong(self):
  675. c = asyncio_client.AsyncClient()
  676. c.state = 'connected'
  677. c.ping_interval = 10
  678. c._send_packet = AsyncMock()
  679. c.queue = mock.MagicMock()
  680. c.queue.put = AsyncMock()
  681. states = [
  682. ('connected', False)
  683. ]
  684. async def fake_wait():
  685. c.state, c.pong_received = states.pop(0)
  686. c.ping_loop_event.wait = fake_wait
  687. _run(c._ping_loop())
  688. self.assertEqual(c.state, 'connected')
  689. c.queue.put.mock.assert_called_once_with(None)
  690. def test_ping_loop_missing_pong_websocket(self):
  691. c = asyncio_client.AsyncClient()
  692. c.state = 'connected'
  693. c.ping_interval = 10
  694. c._send_packet = AsyncMock()
  695. c.queue = mock.MagicMock()
  696. c.queue.put = AsyncMock()
  697. c.ws = mock.MagicMock()
  698. c.ws.close = AsyncMock()
  699. states = [
  700. ('connected', False)
  701. ]
  702. async def fake_wait():
  703. c.state, c.pong_received = states.pop(0)
  704. c.ping_loop_event.wait = fake_wait
  705. _run(c._ping_loop())
  706. self.assertEqual(c.state, 'connected')
  707. c.queue.put.mock.assert_called_once_with(None)
  708. c.ws.close.mock.assert_called_once_with()
  709. def test_read_loop_polling_disconnected(self):
  710. c = asyncio_client.AsyncClient()
  711. c.state = 'disconnected'
  712. c._trigger_event = AsyncMock()
  713. c.write_loop_task = AsyncMock()()
  714. c.ping_loop_task = AsyncMock()()
  715. _run(c._read_loop_polling())
  716. c._trigger_event.mock.assert_not_called()
  717. # should not block
  718. @mock.patch('engineio.client.time.time', return_value=123.456)
  719. def test_read_loop_polling_no_response(self, _time):
  720. c = asyncio_client.AsyncClient()
  721. c.state = 'connected'
  722. c.base_url = 'http://foo'
  723. c.queue = mock.MagicMock()
  724. c.queue.put = AsyncMock()
  725. c._send_request = AsyncMock(return_value=None)
  726. c._trigger_event = AsyncMock()
  727. c.write_loop_task = AsyncMock()()
  728. c.ping_loop_task = AsyncMock()()
  729. _run(c._read_loop_polling())
  730. self.assertEqual(c.state, 'disconnected')
  731. c.queue.put.mock.assert_called_once_with(None)
  732. c._send_request.mock.assert_called_once_with(
  733. 'GET', 'http://foo&t=123.456')
  734. c._trigger_event.mock.assert_called_once_with('disconnect',
  735. run_async=False)
  736. @mock.patch('engineio.client.time.time', return_value=123.456)
  737. def test_read_loop_polling_bad_status(self, _time):
  738. c = asyncio_client.AsyncClient()
  739. c.state = 'connected'
  740. c.base_url = 'http://foo'
  741. c.queue = mock.MagicMock()
  742. c.queue.put = AsyncMock()
  743. c._send_request = AsyncMock()
  744. c._send_request.mock.return_value.status = 400
  745. c.write_loop_task = AsyncMock()()
  746. c.ping_loop_task = AsyncMock()()
  747. _run(c._read_loop_polling())
  748. self.assertEqual(c.state, 'disconnected')
  749. c.queue.put.mock.assert_called_once_with(None)
  750. c._send_request.mock.assert_called_once_with(
  751. 'GET', 'http://foo&t=123.456')
  752. @mock.patch('engineio.client.time.time', return_value=123.456)
  753. def test_read_loop_polling_bad_packet(self, _time):
  754. c = asyncio_client.AsyncClient()
  755. c.state = 'connected'
  756. c.base_url = 'http://foo'
  757. c.queue = mock.MagicMock()
  758. c.queue.put = AsyncMock()
  759. c._send_request = AsyncMock()
  760. c._send_request.mock.return_value.status = 200
  761. c._send_request.mock.return_value.read = AsyncMock(
  762. return_value=b'foo')
  763. c.write_loop_task = AsyncMock()()
  764. c.ping_loop_task = AsyncMock()()
  765. _run(c._read_loop_polling())
  766. self.assertEqual(c.state, 'disconnected')
  767. c.queue.put.mock.assert_called_once_with(None)
  768. c._send_request.mock.assert_called_once_with(
  769. 'GET', 'http://foo&t=123.456')
  770. def test_read_loop_polling(self):
  771. c = asyncio_client.AsyncClient()
  772. c.state = 'connected'
  773. c.base_url = 'http://foo'
  774. c.queue = mock.MagicMock()
  775. c.queue.put = AsyncMock()
  776. c._send_request = AsyncMock()
  777. c._send_request.mock.side_effect = [
  778. mock.MagicMock(status=200, read=AsyncMock(
  779. return_value=payload.Payload(packets=[
  780. packet.Packet(packet.PING),
  781. packet.Packet(packet.NOOP)]).encode())),
  782. None
  783. ]
  784. c.write_loop_task = AsyncMock()()
  785. c.ping_loop_task = AsyncMock()()
  786. c._receive_packet = AsyncMock()
  787. _run(c._read_loop_polling())
  788. self.assertEqual(c.state, 'disconnected')
  789. c.queue.put.mock.assert_called_once_with(None)
  790. self.assertEqual(c._send_request.mock.call_count, 2)
  791. self.assertEqual(c._receive_packet.mock.call_count, 2)
  792. self.assertEqual(
  793. c._receive_packet.mock.call_args_list[0][0][0].encode(), b'2')
  794. self.assertEqual(
  795. c._receive_packet.mock.call_args_list[1][0][0].encode(), b'6')
  796. def test_read_loop_websocket_disconnected(self):
  797. c = asyncio_client.AsyncClient()
  798. c.state = 'disconnected'
  799. c.write_loop_task = AsyncMock()()
  800. c.ping_loop_task = AsyncMock()()
  801. _run(c._read_loop_websocket())
  802. # should not block
  803. def test_read_loop_websocket_no_response(self):
  804. c = asyncio_client.AsyncClient()
  805. c.base_url = 'ws://foo'
  806. c.state = 'connected'
  807. c.queue = mock.MagicMock()
  808. c.queue.put = AsyncMock()
  809. c.ws = mock.MagicMock()
  810. c.ws.recv = AsyncMock(
  811. side_effect=websockets.exceptions.ConnectionClosed(1, 'foo'))
  812. c.write_loop_task = AsyncMock()()
  813. c.ping_loop_task = AsyncMock()()
  814. _run(c._read_loop_websocket())
  815. self.assertEqual(c.state, 'disconnected')
  816. c.queue.put.mock.assert_called_once_with(None)
  817. def test_read_loop_websocket_unexpected_error(self):
  818. c = asyncio_client.AsyncClient()
  819. c.base_url = 'ws://foo'
  820. c.state = 'connected'
  821. c.queue = mock.MagicMock()
  822. c.queue.put = AsyncMock()
  823. c.ws = mock.MagicMock()
  824. c.ws.recv = AsyncMock(side_effect=ValueError)
  825. c.write_loop_task = AsyncMock()()
  826. c.ping_loop_task = AsyncMock()()
  827. _run(c._read_loop_websocket())
  828. self.assertEqual(c.state, 'disconnected')
  829. c.queue.put.mock.assert_called_once_with(None)
  830. def test_read_loop_websocket(self):
  831. c = asyncio_client.AsyncClient()
  832. c.base_url = 'ws://foo'
  833. c.state = 'connected'
  834. c.queue = mock.MagicMock()
  835. c.queue.put = AsyncMock()
  836. c.ws = mock.MagicMock()
  837. c.ws.recv = AsyncMock(side_effect=[
  838. packet.Packet(packet.PING).encode(), ValueError])
  839. c.write_loop_task = AsyncMock()()
  840. c.ping_loop_task = AsyncMock()()
  841. c._receive_packet = AsyncMock()
  842. _run(c._read_loop_websocket())
  843. self.assertEqual(c.state, 'disconnected')
  844. self.assertEqual(
  845. c._receive_packet.mock.call_args_list[0][0][0].encode(), b'2')
  846. c.queue.put.mock.assert_called_once_with(None)
  847. def test_write_loop_disconnected(self):
  848. c = asyncio_client.AsyncClient()
  849. c.state = 'disconnected'
  850. _run(c._write_loop())
  851. # should not block
  852. def test_write_loop_no_packets(self):
  853. c = asyncio_client.AsyncClient()
  854. c.state = 'connected'
  855. c.ping_interval = 1
  856. c.ping_timeout = 2
  857. c.queue = mock.MagicMock()
  858. c.queue.get = AsyncMock(return_value=None)
  859. _run(c._write_loop())
  860. c.queue.task_done.assert_called_once_with()
  861. c.queue.get.mock.assert_called_once_with()
  862. def test_write_loop_empty_queue(self):
  863. c = asyncio_client.AsyncClient()
  864. c.state = 'connected'
  865. c.ping_interval = 1
  866. c.ping_timeout = 2
  867. c.queue = mock.MagicMock()
  868. c.queue.Empty = RuntimeError
  869. c.queue.get = AsyncMock(side_effect=RuntimeError)
  870. _run(c._write_loop())
  871. c.queue.get.mock.assert_called_once_with()
  872. def test_write_loop_polling_one_packet(self):
  873. c = asyncio_client.AsyncClient()
  874. c.base_url = 'http://foo'
  875. c.state = 'connected'
  876. c.ping_interval = 1
  877. c.ping_timeout = 2
  878. c.current_transport = 'polling'
  879. c.queue = mock.MagicMock()
  880. c.queue.Empty = RuntimeError
  881. c.queue.get = AsyncMock(side_effect=[
  882. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  883. RuntimeError
  884. ])
  885. c.queue.get_nowait = mock.MagicMock(side_effect=RuntimeError)
  886. c._send_request = AsyncMock()
  887. c._send_request.mock.return_value.status = 200
  888. _run(c._write_loop())
  889. self.assertEqual(c.queue.task_done.call_count, 1)
  890. p = payload.Payload(
  891. packets=[packet.Packet(packet.MESSAGE, {'foo': 'bar'})])
  892. c._send_request.mock.assert_called_once_with(
  893. 'POST', 'http://foo', body=p.encode(),
  894. headers={'Content-Type': 'application/octet-stream'})
  895. def test_write_loop_polling_three_packets(self):
  896. c = asyncio_client.AsyncClient()
  897. c.base_url = 'http://foo'
  898. c.state = 'connected'
  899. c.ping_interval = 1
  900. c.ping_timeout = 2
  901. c.current_transport = 'polling'
  902. c.queue = mock.MagicMock()
  903. c.queue.Empty = RuntimeError
  904. c.queue.get = AsyncMock(side_effect=[
  905. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  906. RuntimeError
  907. ])
  908. c.queue.get_nowait = mock.MagicMock(side_effect=[
  909. packet.Packet(packet.PING),
  910. packet.Packet(packet.NOOP),
  911. RuntimeError
  912. ])
  913. c._send_request = AsyncMock()
  914. c._send_request.mock.return_value.status = 200
  915. _run(c._write_loop())
  916. self.assertEqual(c.queue.task_done.call_count, 3)
  917. p = payload.Payload(packets=[
  918. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  919. packet.Packet(packet.PING),
  920. packet.Packet(packet.NOOP),
  921. ])
  922. c._send_request.mock.assert_called_once_with(
  923. 'POST', 'http://foo', body=p.encode(),
  924. headers={'Content-Type': 'application/octet-stream'})
  925. def test_write_loop_polling_two_packets_done(self):
  926. c = asyncio_client.AsyncClient()
  927. c.base_url = 'http://foo'
  928. c.state = 'connected'
  929. c.ping_interval = 1
  930. c.ping_timeout = 2
  931. c.current_transport = 'polling'
  932. c.queue = mock.MagicMock()
  933. c.queue.Empty = RuntimeError
  934. c.queue.get = AsyncMock(side_effect=[
  935. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  936. RuntimeError
  937. ])
  938. c.queue.get_nowait = mock.MagicMock(side_effect=[
  939. packet.Packet(packet.PING),
  940. None
  941. ])
  942. c._send_request = AsyncMock()
  943. c._send_request.mock.return_value.status = 200
  944. _run(c._write_loop())
  945. self.assertEqual(c.queue.task_done.call_count, 3)
  946. p = payload.Payload(packets=[
  947. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  948. packet.Packet(packet.PING),
  949. ])
  950. c._send_request.mock.assert_called_once_with(
  951. 'POST', 'http://foo', body=p.encode(),
  952. headers={'Content-Type': 'application/octet-stream'})
  953. self.assertEqual(c.state, 'connected')
  954. def test_write_loop_polling_bad_connection(self):
  955. c = asyncio_client.AsyncClient()
  956. c.base_url = 'http://foo'
  957. c.state = 'connected'
  958. c.ping_interval = 1
  959. c.ping_timeout = 2
  960. c.current_transport = 'polling'
  961. c.queue = mock.MagicMock()
  962. c.queue.Empty = RuntimeError
  963. c.queue.get = AsyncMock(side_effect=[
  964. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  965. ])
  966. c.queue.get_nowait = mock.MagicMock(side_effect=[
  967. RuntimeError
  968. ])
  969. c._send_request = AsyncMock(return_value=None)
  970. _run(c._write_loop())
  971. self.assertEqual(c.queue.task_done.call_count, 1)
  972. p = payload.Payload(
  973. packets=[packet.Packet(packet.MESSAGE, {'foo': 'bar'})])
  974. c._send_request.mock.assert_called_once_with(
  975. 'POST', 'http://foo', body=p.encode(),
  976. headers={'Content-Type': 'application/octet-stream'})
  977. self.assertEqual(c.state, 'connected')
  978. def test_write_loop_polling_bad_status(self):
  979. c = asyncio_client.AsyncClient()
  980. c.base_url = 'http://foo'
  981. c.state = 'connected'
  982. c.ping_interval = 1
  983. c.ping_timeout = 2
  984. c.current_transport = 'polling'
  985. c.queue = mock.MagicMock()
  986. c.queue.Empty = RuntimeError
  987. c.queue.get = AsyncMock(side_effect=[
  988. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  989. ])
  990. c.queue.get_nowait = mock.MagicMock(side_effect=[
  991. RuntimeError
  992. ])
  993. c._send_request = AsyncMock()
  994. c._send_request.mock.return_value.status = 500
  995. _run(c._write_loop())
  996. self.assertEqual(c.queue.task_done.call_count, 1)
  997. p = payload.Payload(
  998. packets=[packet.Packet(packet.MESSAGE, {'foo': 'bar'})])
  999. c._send_request.mock.assert_called_once_with(
  1000. 'POST', 'http://foo', body=p.encode(),
  1001. headers={'Content-Type': 'application/octet-stream'})
  1002. self.assertEqual(c.state, 'disconnected')
  1003. def test_write_loop_websocket_one_packet(self):
  1004. c = asyncio_client.AsyncClient()
  1005. c.state = 'connected'
  1006. c.ping_interval = 1
  1007. c.ping_timeout = 2
  1008. c.current_transport = 'websocket'
  1009. c.queue = mock.MagicMock()
  1010. c.queue.Empty = RuntimeError
  1011. c.queue.get = AsyncMock(side_effect=[
  1012. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  1013. RuntimeError
  1014. ])
  1015. c.queue.get_nowait = mock.MagicMock(side_effect=[
  1016. RuntimeError
  1017. ])
  1018. c.ws = mock.MagicMock()
  1019. c.ws.send = AsyncMock()
  1020. _run(c._write_loop())
  1021. self.assertEqual(c.queue.task_done.call_count, 1)
  1022. self.assertEqual(c.ws.send.mock.call_count, 1)
  1023. c.ws.send.mock.assert_called_once_with('4{"foo":"bar"}')
  1024. def test_write_loop_websocket_three_packets(self):
  1025. c = asyncio_client.AsyncClient()
  1026. c.state = 'connected'
  1027. c.ping_interval = 1
  1028. c.ping_timeout = 2
  1029. c.current_transport = 'websocket'
  1030. c.queue = mock.MagicMock()
  1031. c.queue.Empty = RuntimeError
  1032. c.queue.get = AsyncMock(side_effect=[
  1033. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  1034. RuntimeError
  1035. ])
  1036. c.queue.get_nowait = mock.MagicMock(side_effect=[
  1037. packet.Packet(packet.PING),
  1038. packet.Packet(packet.NOOP),
  1039. RuntimeError
  1040. ])
  1041. c.ws = mock.MagicMock()
  1042. c.ws.send = AsyncMock()
  1043. _run(c._write_loop())
  1044. self.assertEqual(c.queue.task_done.call_count, 3)
  1045. self.assertEqual(c.ws.send.mock.call_count, 3)
  1046. self.assertEqual(c.ws.send.mock.call_args_list[0][0][0],
  1047. '4{"foo":"bar"}')
  1048. self.assertEqual(c.ws.send.mock.call_args_list[1][0][0], '2')
  1049. self.assertEqual(c.ws.send.mock.call_args_list[2][0][0], '6')
  1050. def test_write_loop_websocket_one_packet_binary(self):
  1051. c = asyncio_client.AsyncClient()
  1052. c.state = 'connected'
  1053. c.ping_interval = 1
  1054. c.ping_timeout = 2
  1055. c.current_transport = 'websocket'
  1056. c.queue = mock.MagicMock()
  1057. c.queue.Empty = RuntimeError
  1058. c.queue.get = AsyncMock(side_effect=[
  1059. packet.Packet(packet.MESSAGE, b'foo'),
  1060. RuntimeError
  1061. ])
  1062. c.queue.get_nowait = mock.MagicMock(side_effect=[
  1063. RuntimeError
  1064. ])
  1065. c.ws = mock.MagicMock()
  1066. c.ws.send = AsyncMock()
  1067. _run(c._write_loop())
  1068. self.assertEqual(c.queue.task_done.call_count, 1)
  1069. self.assertEqual(c.ws.send.mock.call_count, 1)
  1070. c.ws.send.mock.assert_called_once_with(b'\x04foo')
  1071. def test_write_loop_websocket_bad_connection(self):
  1072. c = asyncio_client.AsyncClient()
  1073. c.state = 'connected'
  1074. c.ping_interval = 1
  1075. c.ping_timeout = 2
  1076. c.current_transport = 'websocket'
  1077. c.queue = mock.MagicMock()
  1078. c.queue.Empty = RuntimeError
  1079. c.queue.get = AsyncMock(side_effect=[
  1080. packet.Packet(packet.MESSAGE, {'foo': 'bar'}),
  1081. RuntimeError
  1082. ])
  1083. c.queue.get_nowait = mock.MagicMock(side_effect=[
  1084. RuntimeError
  1085. ])
  1086. c.ws = mock.MagicMock()
  1087. c.ws.send = AsyncMock(
  1088. side_effect=websockets.exceptions.ConnectionClosed(1, 'foo'))
  1089. _run(c._write_loop())
  1090. self.assertEqual(c.state, 'connected')