def _eager_bucket_merger_adaptor(iterable, sort_key, bucket_key, buckets):
result = []
merger = streamutils.BucketMerger(sort_key, bucket_key, result.append)
for bucket in buckets:
merger.register(bucket)
for item in iterable:
merger.put(item)
merger.close()
return result
def test_sorting_stable_within_bucket():
emitted = []
merger = streamutils.BucketMerger(lambda x: x[1], lambda x: x[0],
emitted.append)
merger.register("a")
merger.put("a0c")
merger.put("a0a")
merger.put("a0b")
merger.close()
assert emitted == ["a0c", "a0a", "a0b"]
def test_putting_equal_unblocks_all():
emitted = []
merger = streamutils.BucketMerger(lambda x: x[1], lambda x: x[0],
emitted.append)
merger.register("a")
merger.register("b")
merger.put("a0")
assert not emitted
merger.put("b0")
assert set(emitted) == {"a0", "b0"}
def test_unregistering_source_unblocks():
emitted = []
merger = streamutils.BucketMerger(lambda x: x[1], lambda x: x[0],
emitted.append)
merger.register("a")
merger.register("b")
merger.put("a0")
assert not emitted
merger.unregister("b")
assert emitted == ["a0"]
def test_sorting_stable():
emitted = []
merger = streamutils.BucketMerger(lambda x: x[1], lambda x: x[0],
emitted.append)
merger.register("a")
merger.register("b")
merger.put("a0")
merger.put("b1")
merger.put("a1")
merger.close()
assert emitted == ["a0", "b1", "a1"]
def test_putting_greater_unblocks_all_previous():
emitted = []
merger = streamutils.BucketMerger(lambda x: x[1], lambda x: x[0],
emitted.append)
merger.register("a")
merger.register("b")
merger.put("a0")
merger.put("a1")
assert not emitted
merger.put("b2")
assert emitted == ["a0", "a1"]
def test_workflow_and_result_by_fuzzing(num_src, max_event):
random = Random(max_event)
srcs = list(range(num_src))
msgs = list(
_random_interleave(
(_conjure_mminf(
functools.partial(_conjure_msg, src),
random.uniform(1, 10),
random.randint(0, max_event),
random,
) for src in srcs),
random,
))
def get_time(msg):
return msg[1]
def get_src(msg):
return msg[0]
expected = []
actual = []
merger = streamutils.BucketMerger(get_time, get_src, actual.append)
registered_srcs = set()
for src, when in msgs:
if src not in registered_srcs:
registered_srcs.add(src)
merger.register(src)
if actual and when < get_time(actual[-1]):
with pytest.raises(ValueError):
merger.put((src, when))
else:
merger.put((src, when))
expected.append((src, when))
if registered_srcs and random.random() < 0.1:
src = random.choice(list(registered_srcs))
registered_srcs.remove(src)
merger.unregister(src)
merger.close()
# Note that bucket merge is not a stable sort
assert actual == sorted(actual, key=get_time)
assert collections.Counter(actual) == collections.Counter(expected)