239 lines
7.1 KiB
Python
239 lines
7.1 KiB
Python
|
# Copyright (C) PyZMQ Developers
|
||
|
|
# Distributed under the terms of the Modified BSD License.
|
||
|
|
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
|
||
|
|
from pytest import mark
|
||
|
|
|
||
|
|
import zmq
|
||
|
|
|
||
|
|
from zmq.tests import PollZMQTestCase, have_gevent, GreenTest
|
||
|
|
|
||
|
|
def wait():
|
||
|
|
time.sleep(.25)
|
||
|
|
|
||
|
|
|
||
|
|
class TestPoll(PollZMQTestCase):
|
||
|
|
|
||
|
|
Poller = zmq.Poller
|
||
|
|
|
||
|
|
def test_pair(self):
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
|
||
|
|
|
||
|
|
# Sleep to allow sockets to connect.
|
||
|
|
wait()
|
||
|
|
|
||
|
|
poller = self.Poller()
|
||
|
|
poller.register(s1, zmq.POLLIN|zmq.POLLOUT)
|
||
|
|
poller.register(s2, zmq.POLLIN|zmq.POLLOUT)
|
||
|
|
# Poll result should contain both sockets
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
# Now make sure that both are send ready.
|
||
|
|
self.assertEqual(socks[s1], zmq.POLLOUT)
|
||
|
|
self.assertEqual(socks[s2], zmq.POLLOUT)
|
||
|
|
# Now do a send on both, wait and test for zmq.POLLOUT|zmq.POLLIN
|
||
|
|
s1.send(b'msg1')
|
||
|
|
s2.send(b'msg2')
|
||
|
|
wait()
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s1], zmq.POLLOUT|zmq.POLLIN)
|
||
|
|
self.assertEqual(socks[s2], zmq.POLLOUT|zmq.POLLIN)
|
||
|
|
# Make sure that both are in POLLOUT after recv.
|
||
|
|
s1.recv()
|
||
|
|
s2.recv()
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s1], zmq.POLLOUT)
|
||
|
|
self.assertEqual(socks[s2], zmq.POLLOUT)
|
||
|
|
|
||
|
|
poller.unregister(s1)
|
||
|
|
poller.unregister(s2)
|
||
|
|
|
||
|
|
|
||
|
|
def test_reqrep(self):
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.REP, zmq.REQ)
|
||
|
|
|
||
|
|
# Sleep to allow sockets to connect.
|
||
|
|
wait()
|
||
|
|
|
||
|
|
poller = self.Poller()
|
||
|
|
poller.register(s1, zmq.POLLIN|zmq.POLLOUT)
|
||
|
|
poller.register(s2, zmq.POLLIN|zmq.POLLOUT)
|
||
|
|
|
||
|
|
# Make sure that s1 is in state 0 and s2 is in POLLOUT
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(s1 in socks, 0)
|
||
|
|
self.assertEqual(socks[s2], zmq.POLLOUT)
|
||
|
|
|
||
|
|
# Make sure that s2 goes immediately into state 0 after send.
|
||
|
|
s2.send(b'msg1')
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(s2 in socks, 0)
|
||
|
|
|
||
|
|
# Make sure that s1 goes into POLLIN state after a time.sleep().
|
||
|
|
time.sleep(0.5)
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s1], zmq.POLLIN)
|
||
|
|
|
||
|
|
# Make sure that s1 goes into POLLOUT after recv.
|
||
|
|
s1.recv()
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s1], zmq.POLLOUT)
|
||
|
|
|
||
|
|
# Make sure s1 goes into state 0 after send.
|
||
|
|
s1.send(b'msg2')
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(s1 in socks, 0)
|
||
|
|
|
||
|
|
# Wait and then see that s2 is in POLLIN.
|
||
|
|
time.sleep(0.5)
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s2], zmq.POLLIN)
|
||
|
|
|
||
|
|
# Make sure that s2 is in POLLOUT after recv.
|
||
|
|
s2.recv()
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s2], zmq.POLLOUT)
|
||
|
|
|
||
|
|
poller.unregister(s1)
|
||
|
|
poller.unregister(s2)
|
||
|
|
|
||
|
|
def test_no_events(self):
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
|
||
|
|
poller = self.Poller()
|
||
|
|
poller.register(s1, zmq.POLLIN|zmq.POLLOUT)
|
||
|
|
poller.register(s2, 0)
|
||
|
|
self.assertTrue(s1 in poller)
|
||
|
|
self.assertFalse(s2 in poller)
|
||
|
|
poller.register(s1, 0)
|
||
|
|
self.assertFalse(s1 in poller)
|
||
|
|
|
||
|
|
def test_pubsub(self):
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.PUB, zmq.SUB)
|
||
|
|
s2.setsockopt(zmq.SUBSCRIBE, b'')
|
||
|
|
|
||
|
|
# Sleep to allow sockets to connect.
|
||
|
|
wait()
|
||
|
|
|
||
|
|
poller = self.Poller()
|
||
|
|
poller.register(s1, zmq.POLLIN|zmq.POLLOUT)
|
||
|
|
poller.register(s2, zmq.POLLIN)
|
||
|
|
|
||
|
|
# Now make sure that both are send ready.
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s1], zmq.POLLOUT)
|
||
|
|
self.assertEqual(s2 in socks, 0)
|
||
|
|
# Make sure that s1 stays in POLLOUT after a send.
|
||
|
|
s1.send(b'msg1')
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s1], zmq.POLLOUT)
|
||
|
|
|
||
|
|
# Make sure that s2 is POLLIN after waiting.
|
||
|
|
wait()
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(socks[s2], zmq.POLLIN)
|
||
|
|
|
||
|
|
# Make sure that s2 goes into 0 after recv.
|
||
|
|
s2.recv()
|
||
|
|
socks = dict(poller.poll())
|
||
|
|
self.assertEqual(s2 in socks, 0)
|
||
|
|
|
||
|
|
poller.unregister(s1)
|
||
|
|
poller.unregister(s2)
|
||
|
|
|
||
|
|
@mark.skipif(sys.platform.startswith('win'), reason='Windows')
|
||
|
|
def test_raw(self):
|
||
|
|
r, w = os.pipe()
|
||
|
|
r = os.fdopen(r, 'rb')
|
||
|
|
w = os.fdopen(w, 'wb')
|
||
|
|
p = self.Poller()
|
||
|
|
p.register(r, zmq.POLLIN)
|
||
|
|
socks = dict(p.poll(1))
|
||
|
|
assert socks == {}
|
||
|
|
w.write(b'x')
|
||
|
|
w.flush()
|
||
|
|
socks = dict(p.poll(1))
|
||
|
|
assert socks == {r.fileno(): zmq.POLLIN}
|
||
|
|
w.close()
|
||
|
|
r.close()
|
||
|
|
|
||
|
|
def test_timeout(self):
|
||
|
|
"""make sure Poller.poll timeout has the right units (milliseconds)."""
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
|
||
|
|
poller = self.Poller()
|
||
|
|
poller.register(s1, zmq.POLLIN)
|
||
|
|
tic = time.time()
|
||
|
|
evt = poller.poll(.005)
|
||
|
|
toc = time.time()
|
||
|
|
self.assertTrue(toc-tic < 0.1)
|
||
|
|
tic = time.time()
|
||
|
|
evt = poller.poll(5)
|
||
|
|
toc = time.time()
|
||
|
|
self.assertTrue(toc-tic < 0.1)
|
||
|
|
self.assertTrue(toc-tic > .001)
|
||
|
|
tic = time.time()
|
||
|
|
evt = poller.poll(500)
|
||
|
|
toc = time.time()
|
||
|
|
self.assertTrue(toc-tic < 1)
|
||
|
|
self.assertTrue(toc-tic > 0.1)
|
||
|
|
|
||
|
|
class TestSelect(PollZMQTestCase):
|
||
|
|
|
||
|
|
def test_pair(self):
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
|
||
|
|
|
||
|
|
# Sleep to allow sockets to connect.
|
||
|
|
wait()
|
||
|
|
|
||
|
|
rlist, wlist, xlist = zmq.select([s1, s2], [s1, s2], [s1, s2])
|
||
|
|
self.assert_(s1 in wlist)
|
||
|
|
self.assert_(s2 in wlist)
|
||
|
|
self.assert_(s1 not in rlist)
|
||
|
|
self.assert_(s2 not in rlist)
|
||
|
|
|
||
|
|
def test_timeout(self):
|
||
|
|
"""make sure select timeout has the right units (seconds)."""
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
|
||
|
|
tic = time.time()
|
||
|
|
r,w,x = zmq.select([s1,s2],[],[],.005)
|
||
|
|
toc = time.time()
|
||
|
|
self.assertTrue(toc-tic < 1)
|
||
|
|
self.assertTrue(toc-tic > 0.001)
|
||
|
|
tic = time.time()
|
||
|
|
r,w,x = zmq.select([s1,s2],[],[],.25)
|
||
|
|
toc = time.time()
|
||
|
|
self.assertTrue(toc-tic < 1)
|
||
|
|
self.assertTrue(toc-tic > 0.1)
|
||
|
|
|
||
|
|
|
||
|
|
if have_gevent:
|
||
|
|
import gevent
|
||
|
|
from zmq import green as gzmq
|
||
|
|
|
||
|
|
class TestPollGreen(GreenTest, TestPoll):
|
||
|
|
Poller = gzmq.Poller
|
||
|
|
|
||
|
|
def test_wakeup(self):
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
|
||
|
|
poller = self.Poller()
|
||
|
|
poller.register(s2, zmq.POLLIN)
|
||
|
|
|
||
|
|
tic = time.time()
|
||
|
|
r = gevent.spawn(lambda: poller.poll(10000))
|
||
|
|
s = gevent.spawn(lambda: s1.send(b'msg1'))
|
||
|
|
r.join()
|
||
|
|
toc = time.time()
|
||
|
|
self.assertTrue(toc-tic < 1)
|
||
|
|
|
||
|
|
def test_socket_poll(self):
|
||
|
|
s1, s2 = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
|
||
|
|
|
||
|
|
tic = time.time()
|
||
|
|
r = gevent.spawn(lambda: s2.poll(10000))
|
||
|
|
s = gevent.spawn(lambda: s1.send(b'msg1'))
|
||
|
|
r.join()
|
||
|
|
toc = time.time()
|
||
|
|
self.assertTrue(toc-tic < 1)
|
||
|
|
|