forked from goodboy/tractor
1
0
Fork 0
tractor/tests/test_pubsub.py

188 lines
5.2 KiB
Python
Raw Normal View History

2019-01-24 03:35:04 +00:00
import time
2019-01-21 17:31:03 +00:00
from itertools import cycle
import pytest
import trio
import tractor
from async_generator import aclosing
from tractor.testing import tractor_test
def test_type_checks():
with pytest.raises(TypeError) as err:
@tractor.msg.pub
async def no_get_topics(yo):
yield
assert "must define a `get_topics`" in str(err.value)
with pytest.raises(TypeError) as err:
@tractor.msg.pub
def not_async_gen(yo):
pass
assert "must be an async generator function" in str(err.value)
2019-01-21 17:31:03 +00:00
def is_even(i):
return i % 2 == 0
@tractor.msg.pub
async def pubber(get_topics):
2019-01-24 03:35:04 +00:00
ss = tractor.current_actor().statespace
2019-01-21 17:31:03 +00:00
for i in cycle(range(10)):
2019-01-24 03:35:04 +00:00
# ensure topic subscriptions are as expected
ss['get_topics'] = get_topics
2019-01-21 17:31:03 +00:00
yield {'even' if is_even(i) else 'odd': i}
await trio.sleep(0.1)
2019-01-24 03:35:04 +00:00
async def subs(which, pub_actor_name):
2019-01-21 17:31:03 +00:00
if len(which) == 1:
if which[0] == 'even':
pred = is_even
else:
def pred(i):
return not is_even(i)
2019-01-21 17:31:03 +00:00
else:
def pred(i):
return isinstance(i, int)
2019-01-21 17:31:03 +00:00
2019-01-24 03:35:04 +00:00
async with tractor.find_actor(pub_actor_name) as portal:
2019-01-21 17:31:03 +00:00
agen = await portal.run(__name__, 'pubber', topics=which)
async with aclosing(agen) as agen:
async for pkt in agen:
for topic, value in pkt.items():
assert pred(value)
@tractor.msg.pub(tasks=['one', 'two'])
async def multilock_pubber(get_topics):
yield {'doggy': 10}
@pytest.mark.parametrize(
'callwith_expecterror',
[
(pubber, {}, TypeError),
# missing a `topics`
(multilock_pubber, {'ctx': None}, TypeError),
# missing a `task_name`
(multilock_pubber, {'ctx': None, 'topics': ['topic1']}, TypeError),
# should work
(multilock_pubber,
{'ctx': None, 'topics': ['topic1'], 'task_name': 'one'},
None),
],
)
@tractor_test
async def test_required_args(callwith_expecterror):
func, kwargs, err = callwith_expecterror
if err is not None:
with pytest.raises(err):
await func(**kwargs)
else:
async with tractor.open_nursery() as n:
# await func(**kwargs)
portal = await n.run_in_actor(
'sub', multilock_pubber, **kwargs)
async for val in await portal.result():
assert val == {'doggy': 10}
2019-01-24 03:35:04 +00:00
@pytest.mark.parametrize(
'pub_actor',
['streamer', 'arbiter']
)
2019-01-21 17:31:03 +00:00
def test_pubsub_multi_actor_subs(
loglevel,
arb_addr,
2019-01-24 03:35:04 +00:00
pub_actor,
2019-01-21 17:31:03 +00:00
):
2019-01-24 03:35:04 +00:00
"""Try out the neato @pub decorator system.
"""
2019-01-21 17:31:03 +00:00
async def main():
2019-01-24 03:35:04 +00:00
ss = tractor.current_actor().statespace
2019-01-21 17:31:03 +00:00
async with tractor.open_nursery() as n:
2019-01-24 03:35:04 +00:00
name = 'arbiter'
if pub_actor is 'streamer':
# start the publisher as a daemon
master_portal = await n.start_actor(
'streamer',
rpc_module_paths=[__name__],
)
even_portal = await n.run_in_actor(
'evens', subs, which=['even'], pub_actor_name=name)
odd_portal = await n.run_in_actor(
'odds', subs, which=['odd'], pub_actor_name=name)
async with tractor.wait_for_actor('evens'):
# block until 2nd actor is initialized
pass
if pub_actor is 'arbiter':
# wait for publisher task to be spawned in a local RPC task
while not ss.get('get_topics'):
await trio.sleep(0.1)
get_topics = ss.get('get_topics')
assert 'even' in get_topics()
2019-01-21 17:31:03 +00:00
async with tractor.wait_for_actor('odds'):
# block until 2nd actor is initialized
pass
2019-01-24 03:35:04 +00:00
if pub_actor is 'arbiter':
start = time.time()
while 'odd' not in get_topics():
await trio.sleep(0.1)
if time.time() - start > 1:
pytest.fail("odds subscription never arrived?")
2019-01-21 17:31:03 +00:00
# TODO: how to make this work when the arbiter gets
# a portal to itself? Currently this causes a hang
# when the channel server is torn down due to a lingering
# loopback channel
# with trio.move_on_after(1):
# await subs(['even', 'odd'])
# XXX: this would cause infinite
# blocking due to actor never terminating loop
# await even_portal.result()
2019-01-24 03:35:04 +00:00
await trio.sleep(0.5)
2019-01-21 17:31:03 +00:00
await even_portal.cancel_actor()
2019-01-24 03:35:04 +00:00
await trio.sleep(0.5)
if pub_actor is 'arbiter':
assert 'even' not in get_topics()
2019-01-21 17:31:03 +00:00
await odd_portal.cancel_actor()
2019-01-24 03:35:04 +00:00
await trio.sleep(1)
2019-01-21 17:31:03 +00:00
2019-01-24 03:35:04 +00:00
if pub_actor is 'arbiter':
while get_topics():
await trio.sleep(0.1)
if time.time() - start > 1:
pytest.fail("odds subscription never dropped?")
else:
await master_portal.cancel_actor()
2019-01-21 17:31:03 +00:00
tractor.run(
main,
arbiter_addr=arb_addr,
rpc_module_paths=[__name__],
)