intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 39 insertions, 0 deletions

diff --git a/contrib/python/prompt-toolkit/py2/tests/test_utils.py b/contrib/python/prompt-toolkit/py2/tests/test_utils.py
new file mode 100644
index 0000000000..1e3d92cda6
--- /dev/null
+++ b/contrib/python/prompt-toolkit/py2/tests/test_utils.py
@@ -0,0 +1,39 @@
+from __future__ import unicode_literals
+
+from prompt_toolkit.utils import take_using_weights
+
+import itertools
+
+
+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