async def test_sequential(app, results, batch_size):
seq = count()
bulk = 33
keys = list("ABCDEFGHIJ")
actions = app.stream(
"actions",
record=Action,
partition_by="key",
partition_count=3,
)
async def send():
await actions.send(*[Action(key=random.choice(keys), seq=next(seq)) for _ in range(bulk)])
@app.processor(actions, grace_period=1, assignment_sleep=0.3)
async def proc(events):
last_seen = {key: -1 for key in keys}
if batch_size == 1:
async for obj in events.records():
assert obj.seq > last_seen[obj.key]
last_seen[obj.key] = obj.seq
await results.incr()
await results.redis.incr(event_id(obj.seq))
else:
async for batch in events.take(batch_size, within=0.1).records():
for obj in batch:
assert obj.seq > last_seen[obj.key]
last_seen[obj.key] = obj.seq
await results.incr()
await results.redis.incr(event_id(obj.seq))
# Add one worker at a time, then remove one worker at a time, waiting
# for rebalances to complete each iteration.
async with worker(app) as w1:
await wait_running(w1)
async with worker(app) as w2:
await wait_running(w1, w2)
async with worker(app) as w3:
await wait_running(w1, w2, w3)
await send()
# Ensure existing events are processed, then kill workers.
await wait_done(results, count=bulk, delay=4, debug_key=event_id)
await wait_running(w1, w2)
await wait_running(w1)
await send()
# Ensure all events are processed once.
await wait_done(results, count=bulk * 2, delay=8, debug_key=event_id)
async def test_policy_quarantine(app, results):
seq = count()
keys = ["A", "B"]
actions = app.stream(
"actions",
record=Action,
partition_by="key",
partition_count=2,
hasher=lambda key: 0 if key == "A" else 1,
)
# Ensure that events are split between the two partitions.
assert actions.route(keys[0]) != actions.route(keys[1])
@app.processor(actions, exception_policy=ExceptionPolicy.QUARANTINE)
async def proc(events):
async for obj in events.records():
await results.incr()
await results.redis.incr(event_id(obj.seq))
if obj.seq == 55:
raise ValueError("testing-user-error")
async with worker(app):
records = [Action(key="A" if i % 2 == 0 else "B", seq=next(seq)) for i in range(100)]
await actions.send(*records)
# Until seq=55, all events succeed. Then the partition containing key B is
# poisoned, so only half the remaining succeed.
await wait_done(results, count=55 + (45 // 2), debug_key=event_id)
async def test_batch_complex(app, results, want):
orders = app.stream("orders",
record=Order,
partition_by="id",
partition_count=1)
await orders.send(*[Order.random() for _ in range(4)])
@app.processor(orders)
async def proc(events):
if want == "raw":
async for batch in events.take(2, within=0.1):
assert len(batch) == 2
for event in batch:
assert isinstance(event, Event)
await results.incr()
elif want == "records":
async for batch in events.take(2, within=0.1).records():
assert len(batch) == 2
for obj in batch:
assert isinstance(obj.id, int)
assert isinstance(obj.items, list)
await results.incr()
async with worker(app):
await wait_done(results, count=4)
async def test_middleware(app, results, primitive):
class Action(Record, primitive=primitive):
key: str
seq: int
seq = count()
keys = list("ABCDEFGHIJ")
actions = app.stream("actions", record=Action, partition_by="key")
await actions.send(
*[Action(key=random.choice(keys), seq=next(seq)) for _ in range(100)])
dead_letter = DeadLetterMiddleware(source=actions)
@app.processor(actions,
exception_policy=ExceptionPolicy.IGNORE,
middleware=[dead_letter])
async def proc(events):
async for event in events:
obj = actions.deserialize(event.data)
if obj.seq == 11:
raise ValueError("testing-user-error")
@app.processor(dead_letter.sink)
async def dead(events):
async for event in events:
if primitive == True:
assert actions.deserialize(event.data).seq == 11
else:
assert json.loads(
dead_letter.sink.deserialize(event.data).data)["seq"] == 11
await results.incr()
async with worker(app):
await wait_done(results, count=1)
async def test_backpressure(app, results):
"""
Ensure that a single slow partition processor does not block others.
"""
seq = count()
keys = ["A", "B"]
actions = app.stream(
"actions",
record=Action,
partition_by="key",
partition_count=2,
hasher=lambda key: 0 if key == "A" else 1,
)
# Ensure that events are split between the two partitions.
assert actions.route(keys[0]) != actions.route(keys[1])
@app.processor(actions, prefetch_count=8)
async def proc(events):
async for action in events.records():
assert action.key in keys
if action.key == "A":
await asyncio.sleep(10)
await results.incr()
async with worker(app):
records = [
Action(key="A" if i % 2 == 0 else "B", seq=next(seq))
for i in range(100)
]
await actions.send(*records)
# The 50 odd numbered events should be read and processed by partition B,
# even though A is very slow.
await wait_done(results, count=50)
async def test_task(app, results):
@app.task
async def incr():
await results.incr()
async with worker(app) as w:
await wait_started(w)
await wait_done(results, count=1, delay=1)
async def test_timer(app, results):
@app.timer(interval=0.01)
async def incr():
await results.incr()
async with worker(app) as w:
await wait_started(w)
await wait_atleast(results, count=20, delay=1)
async def test_crontab(mocker, app, results):
mocked_seconds_until = mocker.patch("runnel.app.seconds_until")
mocked_seconds_until.return_value = 0.01
@app.crontab("* * * * *")
async def incr():
await results.incr()
async with worker(app) as w:
await wait_started(w)
await wait_atleast(results, count=1, delay=1)
async def test_policy_halt(app, results):
seq = count()
keys = list("ABCDEFGHIJ")
actions = app.stream("actions", record=Action, partition_by="key")
await actions.send(*[Action(key=random.choice(keys), seq=next(seq)) for _ in range(100)])
@app.processor(actions, exception_policy=ExceptionPolicy.HALT, grace_period=0.1)
async def proc(events):
async for obj in events.records():
if obj.seq == 10:
raise ValueError("testing-user-error")
with pytest.raises(ValueError):
async with worker(app):
await anyio.sleep(1)
async def test_policy_ignore(app, results):
seq = count()
keys = list("ABCDEFGHIJ")
actions = app.stream("actions", record=Action, partition_by="key")
await actions.send(*[Action(key=random.choice(keys), seq=next(seq)) for _ in range(100)])
@app.processor(actions, exception_policy=ExceptionPolicy.IGNORE)
async def proc(events):
async for obj in events.records():
await results.incr()
await results.redis.incr(event_id(obj.seq))
if obj.seq == 10:
raise ValueError("testing-user-error")
async with worker(app):
await wait_done(results, count=100, debug_key=event_id)
async def test_builtin(app, results, compressor):
orders = app.stream(
"orders",
record=Order,
partition_by="id",
serializer=JSONSerializer(compressor=compressor),
)
@app.processor(orders)
async def check(events):
async for obj in events.records():
assert isinstance(obj, Order)
await results.incr()
async with worker(app):
await orders.send(*[Order.random() for _ in range(10)])
await wait_done(results, count=10)
async def test_single_primitive(app, results, want):
readings = app.stream("readings", record=Reading, partition_by="id")
await readings.send(*[Reading.random() for _ in range(2)])
@app.processor(readings)
async def proc(events):
if want == "raw":
async for event in events:
assert isinstance(event, Event)
await results.incr()
elif want == "records":
async for obj in events.records():
assert isinstance(obj, Reading)
await results.incr()
async with worker(app):
await wait_done(results, count=2)
async def test_single_complex(app, results, want):
orders = app.stream("orders", record=Order, partition_by="id")
await orders.send(*[Order.random() for _ in range(2)])
@app.processor(orders)
async def proc(events):
if want == "raw":
async for event in events:
assert isinstance(event, Event)
await results.incr()
elif want == "records":
async for obj in events.records():
assert isinstance(obj.id, int)
assert isinstance(obj.items, list)
await results.incr()
async with worker(app):
await wait_done(results, count=2)
async def test_middleware_batches(app, results):
class Action(Record, primitive=True):
key: str
seq: int
seq = count()
keys = list("ABCDEFGHIJ")
batch_size = 10
actions = app.stream("actions",
record=Action,
partition_by="key",
partition_count=1)
await actions.send(
*[Action(key=random.choice(keys), seq=next(seq)) for _ in range(100)])
dead_letter = DeadLetterMiddleware(source=actions)
@app.processor(
actions,
exception_policy=ExceptionPolicy.IGNORE,
middleware=[dead_letter],
prefetch_count=10,
)
async def proc(events):
async for batch in events.take(batch_size, within=0.1):
for event in batch:
obj = actions.deserialize(event.data)
if obj.seq == 11:
raise ValueError("testing-user-error")
# The entire batch has failed, all 10 events sent to the dead-letter queue.
@app.processor(dead_letter.sink)
async def dead(events):
async for event in events:
assert isinstance(actions.deserialize(event.data).seq, int)
await results.incr()
async with worker(app):
await wait_done(results, count=batch_size)
async def test_simple(app, results):
n_events = 10
actions = app.stream(
"actions",
record=Action,
partition_by=lambda a: a.key,
partition_count=2,
)
keys = list("ABCDEFGHIJ")
await actions.send(
*[Action(key=random.choice(keys), seq=i) for i in range(n_events)])
@app.processor(actions)
async def proc(events):
async for action in events.records():
assert action.key in keys
await results.incr()
await results.redis.incr(event_id(action.seq))
async with worker(app) as w:
await wait_running(w)
await wait_done(results, count=n_events, delay=10, debug_key=event_id)
async def test_batch_primitive(app, results, want):
readings = app.stream("readings",
record=Reading,
partition_by="id",
partition_count=1)
await readings.send(*[Reading.random() for _ in range(4)])
@app.processor(readings)
async def proc(events):
if want == "raw":
async for batch in events.take(2, within=0.1):
assert len(batch) == 2
for event in batch:
assert isinstance(event, Event)
await results.incr()
elif want == "records":
async for batch in events.take(2, within=0.1).records():
assert len(batch) == 2
for obj in batch:
assert isinstance(obj, Reading)
await results.incr()
async with worker(app):
await wait_done(results, count=4)