70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
import threading
|
|
from collections import Counter
|
|
from smooth_wrr import SmoothWeightedRoundRobin as S
|
|
|
|
def test_basic_counts():
|
|
s = S([("a", 3), ("b", 2), ("c", 1)])
|
|
seq = [s.next() for _ in range(12)]
|
|
c = Counter(seq)
|
|
assert c == {"a": 6, "b": 4, "c": 2}, c
|
|
# full cycle repeats
|
|
seq2 = [s.next() for _ in range(12)]
|
|
assert seq2 == seq
|
|
|
|
def test_single():
|
|
s = S([("x", 5)])
|
|
assert [s.next() for _ in range(4)] == ["x"]*4
|
|
|
|
def test_smoothness_no_adjacent_same_when_possible():
|
|
s = S([("a", 2), ("b", 1)])
|
|
seq = [s.next() for _ in range(6)]
|
|
# a appears twice as often; never three a in a row
|
|
for i in range(len(seq)-2):
|
|
assert not (seq[i]==seq[i+1]==seq[i+2]=="a"), seq
|
|
|
|
def test_validation():
|
|
try:
|
|
S([])
|
|
raise AssertionError("empty should raise")
|
|
except ValueError: pass
|
|
for bad in [0, -1]:
|
|
try:
|
|
S([("a", bad)])
|
|
raise AssertionError("bad weight should raise")
|
|
except ValueError: pass
|
|
try:
|
|
S([("a", 1.5)])
|
|
raise AssertionError("float weight should raise")
|
|
except TypeError: pass
|
|
try:
|
|
S([("a", True)])
|
|
raise AssertionError("bool weight should raise")
|
|
except TypeError: pass
|
|
|
|
def test_thread_safety_counts():
|
|
s = S([("a", 3), ("b", 2)])
|
|
N = 10000
|
|
out = []
|
|
lock = threading.Lock()
|
|
def worker(k):
|
|
local = [s.next() for _ in range(k)]
|
|
with lock:
|
|
out.extend(local)
|
|
ts = [threading.Thread(target=worker, args=(N//4,)) for _ in range(4)]
|
|
for t in ts: t.start()
|
|
for t in ts: t.join()
|
|
c = Counter(out)
|
|
assert c["a"] == 3*N//5, c
|
|
assert c["b"] == 2*N//5, c
|
|
|
|
def test_reset():
|
|
s = S([("a",1),("b",1)])
|
|
s.next(); s.next()
|
|
s.reset()
|
|
assert s.peek_weights() == [0,0]
|
|
|
|
if __name__ == "__main__":
|
|
test_basic_counts(); test_single(); test_smoothness_no_adjacent_same_when_possible()
|
|
test_validation(); test_thread_safety_counts(); test_reset()
|
|
print("ALL OK")
|