core: fix call_function_test regression

Second time in 3 weeks. So stupid. This time write tests.
pull/175/head
David Wilson 7 years ago
parent 0f29baa077
commit 085b3d21bd

@ -441,13 +441,14 @@ class Receiver(object):
if msg == _DEAD: if msg == _DEAD:
raise ChannelError(ChannelError.local_msg) raise ChannelError(ChannelError.local_msg)
msg.unpickle() # Cause .remote_msg to be thrown.
return msg return msg
def __iter__(self): def __iter__(self):
while True: while True:
try: try:
yield self.get() msg = self.get()
msg.unpickle() # Cause .remote_msg to be thrown.
yield msg
except ChannelError: except ChannelError:
return return

@ -50,6 +50,7 @@ run_test tests/master_test.py
run_test tests/module_finder_test.py run_test tests/module_finder_test.py
run_test tests/nested_test.py run_test tests/nested_test.py
run_test tests/parent_test.py run_test tests/parent_test.py
run_test tests/receiver_test.py
run_test tests/responder_test.py run_test tests/responder_test.py
run_test tests/router_test.py run_test tests/router_test.py
run_test tests/select_test.py run_test tests/select_test.py

@ -0,0 +1,40 @@
import unittest2
import mitogen.core
import testlib
def yield_stuff_then_die(sender):
for x in xrange(5):
sender.put(x)
sender.close()
return 10
class ConstructorTest(testlib.RouterMixin, testlib.TestCase):
def test_handle(self):
recv = mitogen.core.Receiver(self.router)
self.assertTrue(isinstance(recv.handle, int))
self.assertTrue(recv.handle > 100)
self.router.route(
mitogen.core.Message.pickled(
'hi',
dst_id=0,
handle=recv.handle,
)
)
self.assertEquals('hi', recv.get().unpickle())
class IterationTest(testlib.RouterMixin, testlib.TestCase):
def test_dead_stops_iteration(self):
recv = mitogen.core.Receiver(self.router)
fork = self.router.fork()
ret = fork.call_async(yield_stuff_then_die, recv.to_sender())
self.assertEquals(list(range(5)), list(m.unpickle() for m in recv))
self.assertEquals(10, ret.get().unpickle())
if __name__ == '__main__':
unittest2.main()
Loading…
Cancel
Save