summaryrefslogtreecommitdiffstats
path: root/tests/test_utils.py
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2024-05-04 17:35:20 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2024-05-04 17:35:20 +0000
commite106bf94eff07d9a59771d9ccc4406421e18ab64 (patch)
treeedb6545500e39df9c67aa918a6125bffc8ec1aee /tests/test_utils.py
parentInitial commit. (diff)
downloadprompt-toolkit-e106bf94eff07d9a59771d9ccc4406421e18ab64.tar.xz
prompt-toolkit-e106bf94eff07d9a59771d9ccc4406421e18ab64.zip
Adding upstream version 3.0.36.upstream/3.0.36upstream
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'tests/test_utils.py')
-rw-r--r--tests/test_utils.py76
1 files changed, 76 insertions, 0 deletions
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..6a1c179
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,76 @@
+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)