← Back
Editing: test_zmqstream.py
# Copyright (C) PyZMQ Developers # Distributed under the terms of the Modified BSD License. import asyncio from functools import partial from unittest import TestCase import pytest import zmq try: import tornado from tornado import gen, ioloop from zmq.eventloop import zmqstream except ImportError: tornado = None # type: ignore class TestZMQStream(TestCase): def setUp(self): if tornado is None: pytest.skip() self.context = zmq.Context() self.loop = ioloop.IOLoop(make_current=False) if tornado and tornado.version_info < (5,): self.loop.make_current() async def _make_sockets(): self.push = zmqstream.ZMQStream(self.context.socket(zmq.PUSH)) self.pull = zmqstream.ZMQStream(self.context.socket(zmq.PULL)) self.loop.run_sync(_make_sockets) port = self.push.bind_to_random_port('tcp://127.0.0.1') self.pull.connect('tcp://127.0.0.1:%i' % port) self.stream = self.push self._timeout_task = None def tearDown(self): if self._timeout_task: self._timeout_task.cancel() async def wait(): await gen.with_timeout(self.loop.time() + 1, self._timeout_task) self.loop.run_sync(wait) self.loop.close(all_fds=True) self.context.term() def run_until_timeout(self, timeout=50): timed_out = [] async def sleep_timeout(): try: await gen.sleep(timeout) except asyncio.CancelledError: return timed_out[:] = ['timed out'] self.loop.stop() def make_timeout_task(): if tornado.version_info < (5,): self.loop.add_callback(sleep_timeout) else: self._timeout_task = asyncio.ensure_future(sleep_timeout()) self.loop.run_sync(make_timeout_task) self.loop.start() assert not timed_out def test_callable_check(self): """Ensure callable check works (py3k).""" self.stream.on_send(lambda *args: None) self.stream.on_recv(lambda *args: None) self.assertRaises(AssertionError, self.stream.on_recv, 1) self.assertRaises(AssertionError, self.stream.on_send, 1) self.assertRaises(AssertionError, self.stream.on_recv, zmq) async def test_on_recv_basic(self): sent = [b'basic'] self.push.send_multipart(sent) f = asyncio.Future() def callback(msg): assert msg == sent self.loop.stop() self.pull.on_recv(callback) recvd = await asyncio.wait_for(f, timeout=20) assert recvd == sent def test_on_recv_wake(self): sent = [b'wake'] def callback(msg): assert msg == sent self.loop.stop() self.pull.on_recv(callback) self.loop.call_later(1, lambda: self.push.send_multipart(sent)) self.run_until_timeout()
Save File
Cancel