1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78
|
from __future__ import annotations
import itertools
import pytest
from prompt_toolkit.utils import take_using_weights
def test_using_weights():
def take(generator, count):
return list(itertools.islice(generator, 0, count))
# Check distribution.
data = take(take_using_weights(["A", "B", "C"], [5, 10, 20]), 35)
assert data.count("A") == 5
assert data.count("B") == 10
assert data.count("C") == 20
assert data == [
"A",
"B",
"C",
"C",
"B",
"C",
"C",
"A",
"B",
"C",
"C",
"B",
"C",
"C",
"A",
"B",
"C",
"C",
"B",
"C",
"C",
"A",
"B",
"C",
"C",
"B",
"C",
"C",
"A",
"B",
"C",
"C",
"B",
"C",
"C",
]
# Another order.
data = take(take_using_weights(["A", "B", "C"], [20, 10, 5]), 35)
assert data.count("A") == 20
assert data.count("B") == 10
assert data.count("C") == 5
# Bigger numbers.
data = take(take_using_weights(["A", "B", "C"], [20, 10, 5]), 70)
assert data.count("A") == 40
assert data.count("B") == 20
assert data.count("C") == 10
# Negative numbers.
data = take(take_using_weights(["A", "B", "C"], [-20, 10, 0]), 70)
assert data.count("A") == 0
assert data.count("B") == 70
assert data.count("C") == 0
# All zero-weight items.
with pytest.raises(ValueError):
take(take_using_weights(["A", "B", "C"], [0, 0, 0]), 70)
|