async def before_server_start(_sanic, loop):
mpsc = Receiver(loop=loop)
redis = await storage.get_async_redis_pool(loop)
await redis.subscribe(
mpsc.channel('scoreboard'),
mpsc.channel('stolen_flags'),
)
sio.start_background_task(background_task, mpsc)
async def redis_listener():
redis = await aioredis.create_redis('redis://redishost:6379')
receiver = Receiver()
await redis.subscribe(receiver.channel('sockets:notification:message'),
receiver.channel('sockets:notification:message2'))
while await receiver.wait_message():
sender, msg = await receiver.get()
for sid in users:
await sio.emit('notification:message', {'data': str(msg)}, room=sid)
async def listen_redis(self, websocket, channels):
channel = channels.pop(0)
try:
mpsc = Receiver(loop=asyncio.get_event_loop())
await self.sub.subscribe(
mpsc.channel(channel),
*(mpsc.channel(channel) for channel in channels))
async for channel, msg in mpsc.iter():
if websocket.client_state == WebSocketState.CONNECTED:
await websocket.send_bytes(msg)
except:
import traceback
traceback.print_exc()
await self.close_connections()
finally:
logger.info('Connection closed')
async def test_stopped(create_connection, server, caplog):
sub = await create_connection(server.tcp_address)
pub = await create_connection(server.tcp_address)
mpsc = Receiver()
await sub.execute_pubsub("subscribe", mpsc.channel("channel:1"))
assert mpsc.is_active
mpsc.stop()
caplog.clear()
with caplog.at_level("DEBUG", "aioredis"):
await pub.execute("publish", "channel:1", b"Hello")
await asyncio.sleep(0)
assert len(caplog.record_tuples) == 1
# Receiver must have 1 EndOfStream message
message = (
"Pub/Sub listener message after stop: "
"sender: <_Sender name:b'channel:1', is_pattern:False, receiver:"
"<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'")
assert caplog.record_tuples == [
("aioredis", logging.WARNING, message),
]
# assert (await mpsc.get()) is None
with pytest.raises(ChannelClosedError):
await mpsc.get()
res = await mpsc.wait_message()
assert res is False
async def test_pubsub_receiver_iter(create_redis, server, event_loop):
sub = await create_redis(server.tcp_address)
pub = await create_redis(server.tcp_address)
mpsc = Receiver()
async def coro(mpsc):
lst = []
async for msg in mpsc.iter():
lst.append(msg)
return lst
tsk = asyncio.ensure_future(coro(mpsc))
(snd1, ) = await sub.subscribe(mpsc.channel("chan:1"))
(snd2, ) = await sub.subscribe(mpsc.channel("chan:2"))
(snd3, ) = await sub.psubscribe(mpsc.pattern("chan:*"))
subscribers = await pub.publish_json("chan:1", {"Hello": "World"})
assert subscribers > 1
subscribers = await pub.publish_json("chan:2", ["message"])
assert subscribers > 1
event_loop.call_later(0, mpsc.stop)
await asyncio.sleep(0.01)
assert await tsk == [
(snd1, b'{"Hello": "World"}'),
(snd3, (b"chan:1", b'{"Hello": "World"}')),
(snd2, b'["message"]'),
(snd3, (b"chan:2", b'["message"]')),
]
assert not mpsc.is_active
async def test_pubsub_receiver_iter(create_redis, server, loop):
sub = await create_redis(server.tcp_address, loop=loop)
pub = await create_redis(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
async def coro(mpsc):
lst = []
async for msg in mpsc.iter():
lst.append(msg)
return lst
tsk = asyncio.ensure_future(coro(mpsc), loop=loop)
snd1, = await sub.subscribe(mpsc.channel('chan:1'))
snd2, = await sub.subscribe(mpsc.channel('chan:2'))
snd3, = await sub.psubscribe(mpsc.pattern('chan:*'))
subscribers = await pub.publish_json('chan:1', {'Hello': 'World'})
assert subscribers > 1
subscribers = await pub.publish_json('chan:2', ['message'])
assert subscribers > 1
loop.call_later(0, mpsc.stop)
# await asyncio.sleep(0, loop=loop)
assert await tsk == [
(snd1, b'{"Hello": "World"}'),
(snd3, (b'chan:1', b'{"Hello": "World"}')),
(snd2, b'["message"]'),
(snd3, (b'chan:2', b'["message"]')),
]
assert not mpsc.is_active
class Messager:
def __init__(self, inbound, outbound, loop):
self.loop = loop
self.conn = None
self.inbound = inbound
self.outbound = outbound
self.receiver = None
self.task = None
self.replies = dict()
self.expected = set()
async def initialize(self):
self.conn = await aioredis.create_redis_pool(("localhost", 6379),
encoding="utf-8",
maxsize=2)
self.receiver = Receiver(loop=self.loop)
await self.conn.subscribe(self.receiver.channel(self.inbound))
self.task = self.loop.create_task(self.fetcher())
print(
f'Redis connection established, listening on {self.inbound}, sending on {self.outbound}'
) # FIXME - propper logging
async def terminate(self):
# terminate channels and disconnect from redis
self.conn.unsubscribe(self.inbound)
self.task.cancel()
self.receiver.stop()
self.conn.close()
await self.conn.wait_closed()
async def fetcher(self):
async for sender, message in self.receiver.iter(encoding='utf-8',
decoder=json.loads):
channel = sender.name.decode()
if channel == self.inbound:
uid = message["uid"]
if uid not in self.expected:
print("Unexpected message!")
print(message)
else:
self.expected.remove(uid)
self.replies[uid] = message
async def get_reply(self, data):
try:
return (await asyncio.wait_for(self._get_reply(data), 10))["reply"]
except TimeoutError:
raise Redisception("No reply received from the bot!")
async def _get_reply(self, data):
uid = str(uuid.uuid4())
self.expected.add(uid)
data["uid"] = uid
await self.conn.publish_json(self.outbound, data)
while uid not in self.replies:
await asyncio.sleep(0.1)
reply = self.replies[uid]
del self.replies[uid]
return reply
async def _update_log_info(self):
def process_log_batch(log_batch):
ip = log_batch["ip"]
pid = str(log_batch["pid"])
if pid != "autoscaler":
logs_for_ip = dict(DataSource.ip_and_pid_to_logs.get(ip, {}))
logs_for_pid = list(logs_for_ip.get(pid, []))
logs_for_pid.extend(log_batch["lines"])
# Only cache upto MAX_LOGS_TO_CACHE
logs_length = len(logs_for_pid)
if logs_length > MAX_LOGS_TO_CACHE * LOG_PRUNE_THREASHOLD:
offset = logs_length - MAX_LOGS_TO_CACHE
del logs_for_pid[:offset]
logs_for_ip[pid] = logs_for_pid
DataSource.ip_and_pid_to_logs[ip] = logs_for_ip
logger.info(f"Received a log for {ip} and {pid}")
aioredis_client = self._dashboard_head.aioredis_client
receiver = Receiver()
channel = receiver.channel(gcs_utils.LOG_FILE_CHANNEL)
await aioredis_client.subscribe(channel)
logger.info("Subscribed to %s", channel)
async for sender, msg in receiver.iter():
try:
data = json.loads(ray._private.utils.decode(msg))
data["pid"] = str(data["pid"])
process_log_batch(data)
except Exception:
logger.exception("Error receiving log from Redis.")
async def test_stopped(create_connection, server, loop):
sub = await create_connection(server.tcp_address, loop=loop)
pub = await create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'))
assert mpsc.is_active
mpsc.stop()
with logs('aioredis', 'DEBUG') as cm:
await pub.execute('publish', 'channel:1', b'Hello')
await asyncio.sleep(0, loop=loop)
assert len(cm.output) == 1
# Receiver must have 1 EndOfStream message
warn_messaege = (
"WARNING:aioredis:Pub/Sub listener message after stop: "
"sender: <_Sender name:b'channel:1', is_pattern:False, receiver:"
"<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'"
)
assert cm.output == [warn_messaege]
# assert (await mpsc.get()) is None
with pytest.raises(ChannelClosedError):
await mpsc.get()
res = await mpsc.wait_message()
assert res is False
async def test_stopped(create_connection, server, loop):
sub = await create_connection(server.tcp_address, loop=loop)
pub = await create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'))
assert mpsc.is_active
mpsc.stop()
with logs('aioredis', 'DEBUG') as cm:
await pub.execute('publish', 'channel:1', b'Hello')
await asyncio.sleep(0, loop=loop)
assert len(cm.output) == 1
# Receiver must have 1 EndOfStream message
warn_messaege = (
"WARNING:aioredis:Pub/Sub listener message after stop: "
"sender: <_Sender name:b'channel:1', is_pattern:False, receiver:"
"<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'")
assert cm.output == [warn_messaege]
# assert (await mpsc.get()) is None
with pytest.raises(ChannelClosedError):
await mpsc.get()
res = await mpsc.wait_message()
assert res is False
async def test_pubsub_receiver_iter(create_redis, server, loop):
sub = await create_redis(server.tcp_address, loop=loop)
pub = await create_redis(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
async def coro(mpsc):
lst = []
async for msg in mpsc.iter():
lst.append(msg)
return lst
tsk = asyncio.ensure_future(coro(mpsc), loop=loop)
snd1, = await sub.subscribe(mpsc.channel('chan:1'))
snd2, = await sub.subscribe(mpsc.channel('chan:2'))
snd3, = await sub.psubscribe(mpsc.pattern('chan:*'))
subscribers = await pub.publish_json('chan:1', {'Hello': 'World'})
assert subscribers > 1
subscribers = await pub.publish_json('chan:2', ['message'])
assert subscribers > 1
loop.call_later(0, mpsc.stop)
# await asyncio.sleep(0, loop=loop)
assert await tsk == [
(snd1, b'{"Hello": "World"}'),
(snd3, (b'chan:1', b'{"Hello": "World"}')),
(snd2, b'["message"]'),
(snd3, (b'chan:2', b'["message"]')),
]
assert not mpsc.is_active
async def producer(self):
channel_names = self.channel_names or []
channel_patterns = self.channel_patterns or []
if not channel_names and not channel_patterns:
return
mpsc = Receiver()
if channel_names:
channels = [mpsc.channel(c) for c in channel_names]
await self.redis.subscribe(*channels)
if channel_patterns:
tasks = set()
for p in channel_patterns:
tasks.add(self.redis.psubscribe(mpsc.pattern(p)))
if tasks:
await asyncio.wait(tasks)
try:
await self.receiver_reader(mpsc)
finally:
if channel_names:
await self.redis.unsubscribe(*channel_names)
if channel_patterns:
tasks = set()
for p in channel_patterns:
tasks.add(self.redis.punsubscribe(p))
await asyncio.wait(tasks)
mpsc.stop()
async def redis_relay(websocket):
conn = await aioredis.create_connection(("localhost", 6379))
receiver = Receiver()
conn.execute_pubsub("subscribe", receiver.channel("marketupdates"))
while await receiver.wait_message():
*_, message = await receiver.get()
await websocket.send(message.decode())
async def receive(self):
print('Receive from redis')
sub = await aioredis.create_connection(('localhost', 6379))
receiver = Receiver()
sub.execute_pubsub('subscribe', receiver.channel(self.channel))
while (await receiver.wait_message()):
msg = await receiver.get()
print("Got Message:", msg)
async def subscribe(self):
if not self._subscribe:
redis = await self.redis
# self._subscribe = await redis.subscribe('ws')
r = Receiver(loop=self.loop)
for ch in self._channels:
await redis.subscribe(r.channel(ch))
self._subscribe = r
return self._subscribe
async def subscribe_to_channel(loop, redis_pool):
mpsc = Receiver(loop=loop)
logger.info("Aquiring redis connection from pool")
connection = await redis_pool.acquire()
logger.info("Subscribing to redis channel: {}".format(websocket_channel))
await connection.execute_pubsub("subscribe",
mpsc.channel(websocket_channel))
return mpsc
async def test_unsubscribe(create_connection, server):
sub = await create_connection(server.tcp_address)
pub = await create_connection(server.tcp_address)
mpsc = Receiver()
await sub.execute_pubsub(
"subscribe", mpsc.channel("channel:1"), mpsc.channel("channel:3")
)
res = await pub.execute("publish", "channel:3", "Hello world")
assert res == 1
res = await pub.execute("publish", "channel:1", "Hello world")
assert res == 1
assert mpsc.is_active
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b"channel:3"
assert not ch.is_pattern
assert msg == b"Hello world"
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b"channel:1"
assert not ch.is_pattern
assert msg == b"Hello world"
await sub.execute_pubsub("unsubscribe", "channel:1")
assert mpsc.is_active
res = await pub.execute("publish", "channel:3", "message")
assert res == 1
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b"channel:3"
assert not ch.is_pattern
assert msg == b"message"
waiter = asyncio.ensure_future(mpsc.get())
await sub.execute_pubsub("unsubscribe", "channel:3")
assert not mpsc.is_active
assert await waiter is None
async def _event_listener(self, channel):
async for _ in AsyncCirculator():
async with self._redis_pool.get_client() as cache:
receiver = Receiver()
await cache.subscribe(receiver.channel(channel))
async for channel, message in receiver.iter():
await self._event_assigner(channel, message)
async def test_unsubscribe(create_connection, server, loop):
sub = await create_connection(server.tcp_address, loop=loop)
pub = await create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'),
mpsc.channel('channel:3'))
res = await pub.execute("publish", "channel:3", "Hello world")
assert res == 1
res = await pub.execute("publish", "channel:1", "Hello world")
assert res == 1
assert mpsc.is_active
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b'channel:3'
assert not ch.is_pattern
assert msg == b"Hello world"
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b'channel:1'
assert not ch.is_pattern
assert msg == b"Hello world"
await sub.execute_pubsub('unsubscribe', 'channel:1')
assert mpsc.is_active
res = await pub.execute("publish", "channel:3", "message")
assert res == 1
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b'channel:3'
assert not ch.is_pattern
assert msg == b"message"
await sub.execute_pubsub('unsubscribe', 'channel:3')
assert not mpsc.is_active
res = await mpsc.get()
assert res is None
async def main():
await redis.connect()
mpsc = Receiver(loop=asyncio.get_event_loop())
asyncio.ensure_future(reader(mpsc))
await redis._redis.subscribe(mpsc.channel('channel:1'))
while True:
try:
await asyncio.sleep(10)
print('hearbeat', flush=True)
except Exception as ex:
print(ex, flush=True)
mpsc.stop()
await redis.disconnect()
async def test_subscriptions(create_connection, server):
sub = await create_connection(server.tcp_address)
pub = await create_connection(server.tcp_address)
mpsc = Receiver()
await sub.execute_pubsub("subscribe", mpsc.channel("channel:1"),
mpsc.channel("channel:3"))
res = await pub.execute("publish", "channel:3", "Hello world")
assert res == 1
res = await pub.execute("publish", "channel:1", "Hello world")
assert res == 1
assert mpsc.is_active
ch, msg = await mpsc.get()
assert ch.name == b"channel:3"
assert not ch.is_pattern
assert msg == b"Hello world"
ch, msg = await mpsc.get()
assert ch.name == b"channel:1"
assert not ch.is_pattern
assert msg == b"Hello world"
async def subscribe_chat(self):
redis_for_pubsub = await aioredis.create_redis(self.redis.address, db=self.redis.db)
receiver = Receiver()
channel = receiver.channel(CHAT_CHANNEL)
await redis_for_pubsub.subscribe(channel)
async def _get():
return (await receiver.get())[1]
yield channel, _get
await redis_for_pubsub.unsubscribe(CHAT_CHANNEL)
async def test_subscriptions(create_connection, server, loop):
sub = await create_connection(server.tcp_address, loop=loop)
pub = await create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'),
mpsc.channel('channel:3'))
res = await pub.execute("publish", "channel:3", "Hello world")
assert res == 1
res = await pub.execute("publish", "channel:1", "Hello world")
assert res == 1
assert mpsc.is_active
ch, msg = await mpsc.get()
assert ch.name == b'channel:3'
assert not ch.is_pattern
assert msg == b"Hello world"
ch, msg = await mpsc.get()
assert ch.name == b'channel:1'
assert not ch.is_pattern
assert msg == b"Hello world"
async def test_decode_message_error(loop):
mpsc = Receiver(loop)
ch = mpsc.channel('channel:1')
ch.put_nowait(b'{"hello": "world"}')
unexpected = (mock.ANY, {'hello': 'world'})
with pytest.raises(TypeError):
assert (await mpsc.get(decoder=json.loads)) == unexpected
ch = mpsc.pattern('*')
ch.put_nowait((b'channel', b'{"hello": "world"}'))
unexpected = (mock.ANY, b'channel', {'hello': 'world'})
with pytest.raises(TypeError):
assert (await mpsc.get(decoder=json.loads)) == unexpected
async def test_decode_message_error():
mpsc = Receiver()
ch = mpsc.channel("channel:1")
ch.put_nowait(b'{"hello": "world"}')
unexpected = (mock.ANY, {"hello": "world"})
with pytest.raises(TypeError):
assert (await mpsc.get(decoder=json.loads)) == unexpected
ch = mpsc.pattern("*")
ch.put_nowait((b"channel", b'{"hello": "world"}'))
unexpected = (mock.ANY, b"channel", {"hello": "world"})
with pytest.raises(TypeError):
assert (await mpsc.get(decoder=json.loads)) == unexpected
async def test_decode_message_error(loop):
mpsc = Receiver(loop)
ch = mpsc.channel('channel:1')
ch.put_nowait(b'{"hello": "world"}')
unexpected = (mock.ANY, {'hello': 'world'})
with pytest.raises(TypeError):
assert (await mpsc.get(decoder=json.loads)) == unexpected
ch = mpsc.pattern('*')
ch.put_nowait((b'channel', b'{"hello": "world"}'))
unexpected = (mock.ANY, b'channel', {'hello': 'world'})
with pytest.raises(TypeError):
assert (await mpsc.get(decoder=json.loads)) == unexpected
def test_listener_channel(loop):
mpsc = Receiver(loop=loop)
assert not mpsc.is_active
ch_a = mpsc.channel("channel:1")
assert isinstance(ch_a, AbcChannel)
assert mpsc.is_active
ch_b = mpsc.channel('channel:1')
assert ch_a is ch_b
assert ch_a.name == ch_b.name
assert ch_a.is_pattern == ch_b.is_pattern
assert mpsc.is_active
# remember id; drop refs to objects and create new one;
ch_a.close()
assert not ch_a.is_active
assert not mpsc.is_active
ch = mpsc.channel("channel:1")
assert ch is not ch_a
assert dict(mpsc.channels) == {b'channel:1': ch}
assert dict(mpsc.patterns) == {}
def test_listener_channel(loop):
mpsc = Receiver(loop=loop)
assert not mpsc.is_active
ch_a = mpsc.channel("channel:1")
assert isinstance(ch_a, AbcChannel)
assert mpsc.is_active
ch_b = mpsc.channel('channel:1')
assert ch_a is ch_b
assert ch_a.name == ch_b.name
assert ch_a.is_pattern == ch_b.is_pattern
assert mpsc.is_active
# remember id; drop refs to objects and create new one;
ch_a.close()
assert not ch_a.is_active
assert not mpsc.is_active
ch = mpsc.channel("channel:1")
assert ch is not ch_a
assert dict(mpsc.channels) == {b'channel:1': ch}
assert dict(mpsc.patterns) == {}
async def test_pubsub_receiver_call_stop_with_empty_queue(
create_redis, server, loop):
sub = await create_redis(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
# FIXME: currently at least one subscriber is needed
snd1, = await sub.subscribe(mpsc.channel('chan:1'))
now = loop.time()
loop.call_later(.5, mpsc.stop)
async for i in mpsc.iter(): # noqa (flake8 bug with async for)
assert False, "StopAsyncIteration not raised"
dt = loop.time() - now
assert dt <= 1.5
assert not mpsc.is_active
async def test_pubsub_receiver_call_stop_with_empty_queue(
create_redis, server, loop):
sub = await create_redis(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
# FIXME: currently at least one subscriber is needed
snd1, = await sub.subscribe(mpsc.channel('chan:1'))
now = loop.time()
loop.call_later(.5, mpsc.stop)
async for i in mpsc.iter(): # noqa (flake8 bug with async for)
assert False, "StopAsyncIteration not raised"
dt = loop.time() - now
assert dt <= 1.5
assert not mpsc.is_active
async def subscribe_handler(_request, ws):
data = await ws.recv()
try:
decoded_data = ujson.decode(data)
except ValueError:
await ws.send(ujson.dumps({'error': 'invalid json data'}))
return
token = decoded_data.get('token', '')
loop = asyncio.get_event_loop()
redis = await storage.get_async_redis(loop)
mpsc = Receiver(loop=loop)
await redis.subscribe(mpsc.channel(f'updates:{token}'))
async for channel, msg in mpsc.iter():
await ws.send(msg.decode())
async def test_wait_message(create_connection, server, loop):
sub = await create_connection(server.tcp_address, loop=loop)
pub = await create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'))
fut = asyncio.ensure_future(mpsc.wait_message(), loop=loop)
assert not fut.done()
await asyncio.sleep(0, loop=loop)
assert not fut.done()
await pub.execute('publish', 'channel:1', 'hello')
await asyncio.sleep(0, loop=loop) # read in connection
await asyncio.sleep(0, loop=loop) # call Future.set_result
assert fut.done()
res = await fut
assert res is True
def test_wait_message(create_connection, server, loop):
sub = yield from create_connection(server.tcp_address, loop=loop)
pub = yield from create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
yield from sub.execute_pubsub('subscribe', mpsc.channel('channel:1'))
fut = async_task(mpsc.wait_message(), loop=loop)
assert not fut.done()
yield from asyncio.sleep(0, loop=loop)
assert not fut.done()
yield from pub.execute('publish', 'channel:1', 'hello')
yield from asyncio.sleep(0, loop=loop) # read in connection
yield from asyncio.sleep(0, loop=loop) # call Future.set_result
assert fut.done()
res = yield from fut
assert res is True
async def test_wait_message(create_connection, server, loop):
sub = await create_connection(server.tcp_address, loop=loop)
pub = await create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'))
fut = asyncio.ensure_future(mpsc.wait_message(), loop=loop)
assert not fut.done()
await asyncio.sleep(0, loop=loop)
assert not fut.done()
await pub.execute('publish', 'channel:1', 'hello')
await asyncio.sleep(0, loop=loop) # read in connection
await asyncio.sleep(0, loop=loop) # call Future.set_result
assert fut.done()
res = await fut
assert res is True
async def work(self, app):
"""
Start work, fetch tasks
:param app: LRS object
"""
self.app = app
redis_pool = await aioredis.create_pool(
(app.redis_host, app.redis_port),
minsize=5,
maxsize=10,
password=app.redis_password or None,
db=app.redis_db,
loop=app.loop)
self.redis = await aioredis.Redis(redis_pool)
mpsc = Receiver(loop=app.loop)
asyncio.ensure_future(self.reader(mpsc), loop=app.loop)
await self.redis.subscribe(mpsc.channel('crawl:finish'))
async def test_wait_message(create_connection, server):
sub = await create_connection(server.tcp_address)
pub = await create_connection(server.tcp_address)
mpsc = Receiver()
await sub.execute_pubsub("subscribe", mpsc.channel("channel:1"))
fut = asyncio.ensure_future(mpsc.wait_message())
assert not fut.done()
await asyncio.sleep(0)
assert not fut.done()
await pub.execute("publish", "channel:1", "hello")
await asyncio.sleep(0) # read in connection
await asyncio.sleep(0) # call Future.set_result
assert fut.done()
res = await fut
assert res is True
async def test_decode_message(loop):
mpsc = Receiver(loop)
ch = mpsc.channel('channel:1')
ch.put_nowait(b'Some data')
res = await mpsc.get(encoding='utf-8')
assert isinstance(res[0], _Sender)
assert res[1] == 'Some data'
ch.put_nowait('{"hello": "world"}')
res = await mpsc.get(decoder=json.loads)
assert isinstance(res[0], _Sender)
assert res[1] == {'hello': 'world'}
ch.put_nowait(b'{"hello": "world"}')
res = await mpsc.get(encoding='utf-8', decoder=json.loads)
assert isinstance(res[0], _Sender)
assert res[1] == {'hello': 'world'}
async def test_unsubscribe(create_connection, server, loop):
sub = await create_connection(server.tcp_address, loop=loop)
pub = await create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
await sub.execute_pubsub('subscribe',
mpsc.channel('channel:1'),
mpsc.channel('channel:3'))
res = await pub.execute("publish", "channel:3", "Hello world")
assert res == 1
res = await pub.execute("publish", "channel:1", "Hello world")
assert res == 1
assert mpsc.is_active
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b'channel:3'
assert not ch.is_pattern
assert msg == b"Hello world"
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b'channel:1'
assert not ch.is_pattern
assert msg == b"Hello world"
await sub.execute_pubsub('unsubscribe', 'channel:1')
assert mpsc.is_active
res = await pub.execute("publish", "channel:3", "message")
assert res == 1
assert (await mpsc.wait_message()) is True
ch, msg = await mpsc.get()
assert ch.name == b'channel:3'
assert not ch.is_pattern
assert msg == b"message"
waiter = asyncio.ensure_future(mpsc.get(), loop=loop)
await sub.execute_pubsub('unsubscribe', 'channel:3')
assert not mpsc.is_active
assert await waiter is None
async def test_subscriptions(create_connection, server, loop):
sub = await create_connection(server.tcp_address, loop=loop)
pub = await create_connection(server.tcp_address, loop=loop)
mpsc = Receiver(loop=loop)
await sub.execute_pubsub('subscribe',
mpsc.channel('channel:1'),
mpsc.channel('channel:3'))
res = await pub.execute("publish", "channel:3", "Hello world")
assert res == 1
res = await pub.execute("publish", "channel:1", "Hello world")
assert res == 1
assert mpsc.is_active
ch, msg = await mpsc.get()
assert ch.name == b'channel:3'
assert not ch.is_pattern
assert msg == b"Hello world"
ch, msg = await mpsc.get()
assert ch.name == b'channel:1'
assert not ch.is_pattern
assert msg == b"Hello world"
async def test_pubsub_receiver_stop_on_disconnect(create_redis, server, loop):
pub = await create_redis(server.tcp_address, loop=loop)
sub = await create_redis(server.tcp_address, loop=loop)
sub_name = 'sub-{:X}'.format(id(sub))
await sub.client_setname(sub_name)
for sub_info in await pub.client_list():
if sub_info.name == sub_name:
break
assert sub_info.name == sub_name
mpsc = Receiver(loop=loop)
await sub.subscribe(mpsc.channel('channel:1'))
await sub.subscribe(mpsc.channel('channel:2'))
await sub.psubscribe(mpsc.pattern('channel:*'))
q = asyncio.Queue(loop=loop)
EOF = object()
async def reader():
async for ch, msg in mpsc.iter(encoding='utf-8'):
await q.put((ch.name, msg))
await q.put(EOF)
tsk = asyncio.ensure_future(reader(), loop=loop)
await pub.publish_json('channel:1', ['hello'])
await pub.publish_json('channel:2', ['hello'])
# receive all messages
assert await q.get() == (b'channel:1', '["hello"]')
assert await q.get() == (b'channel:*', (b'channel:1', '["hello"]'))
assert await q.get() == (b'channel:2', '["hello"]')
assert await q.get() == (b'channel:*', (b'channel:2', '["hello"]'))
# XXX: need to implement `client kill`
assert await pub.execute('client', 'kill', sub_info.addr) in (b'OK', 1)
await asyncio.wait_for(tsk, timeout=1, loop=loop)
assert await q.get() is EOF
