diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-15 16:35:31 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-04-15 16:35:31 +0000 |
commit | 4f1a3b5f9ad05aa7b08715d48909a2b06ee2fcb1 (patch) | |
tree | e5dee7be2f0d963da4faad6517278d03783e3adc /tests/test_utils.py | |
parent | Initial commit. (diff) | |
download | prompt-toolkit-7636d8a881b219a91fe81306e0a97e2eb8be010b.tar.xz prompt-toolkit-7636d8a881b219a91fe81306e0a97e2eb8be010b.zip |
Adding upstream version 3.0.43.upstream/3.0.43
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'tests/test_utils.py')
-rw-r--r-- | tests/test_utils.py | 78 |
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) |