1 files changed, 78 insertions, 0 deletions
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..9d4c808
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,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)