From 195571b5551012bd8315359df3425290382c1636 Mon Sep 17 00:00:00 2001
From: Kyle Marsh <kyle.marsh@dreamhost.com>
Date: Wed, 10 Aug 2011 13:26:00 -0700
Subject: [PATCH] S3 Fuzzer: support weights for choices

---
 s3tests/functional/test_fuzzer.py | 27 +++++++++++++++++++++++++++
 s3tests/fuzz_headers.py           | 19 +++++++++++++++++--
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/s3tests/functional/test_fuzzer.py b/s3tests/functional/test_fuzzer.py
index 031e59f..089d52b 100644
--- a/s3tests/functional/test_fuzzer.py
+++ b/s3tests/functional/test_fuzzer.py
@@ -52,6 +52,14 @@ def build_graph():
         },
         'choices': ['leaf']
     }
+    graph['weighted_choices'] = {
+        'set': {},
+        'choices': [
+            'foo',
+            '2 bar',
+            '1 baz'
+        ]
+    }
     return graph
 
 
@@ -149,3 +157,22 @@ def test_expand_decision():
     eq(request['randkey'], 'value-NI$;92@H/0I') #FIXME: again, how to handle the pseudorandom content?
     assert_raises(KeyError, lambda x: decision[x], 'key3')
 
+def test_weighted_choices():
+    graph = build_graph()
+    prng = random.Random(1)
+
+    choices_made = {}
+    for _ in xrange(1000):
+        choice = make_choice(graph['weighted_choices']['choices'], prng)
+        if choices_made.has_key(choice):
+            choices_made[choice] += 1
+        else:
+            choices_made[choice] = 1
+
+    foo_percentage = choices_made['foo'] / 1000.0
+    bar_percentage = choices_made['bar'] / 1000.0
+    baz_percentage = choices_made['baz'] / 1000.0
+    nose.tools.assert_almost_equal(foo_percentage, 0.25, 1)
+    nose.tools.assert_almost_equal(bar_percentage, 0.50, 1)
+    nose.tools.assert_almost_equal(baz_percentage, 0.25, 1)
+
diff --git a/s3tests/fuzz_headers.py b/s3tests/fuzz_headers.py
index 54095b6..91356f8 100644
--- a/s3tests/fuzz_headers.py
+++ b/s3tests/fuzz_headers.py
@@ -26,8 +26,7 @@ def descend_graph(decision_graph, node_name, prng):
     node = decision_graph[node_name]
 
     try:
-        #TODO: Give weights to each choice
-        choice = prng.choice(node['choices'])
+        choice = make_choice(node['choices'], prng)
         decision = descend_graph(decision_graph, choice, prng)
     except IndexError:
         decision = {}
@@ -40,6 +39,22 @@ def descend_graph(decision_graph, node_name, prng):
     return decision
 
 
+def make_choice(choices, prng):
+    weighted_choices = []
+    for option in choices:
+        fields = option.split(None, 1)
+        if len(fields) == 1:
+            weight = 1
+            value = fields[0]
+        else:
+            weight = int(fields[0])
+            value = fields[1]
+        for _ in xrange(weight):
+            weighted_choices.append(value)
+
+    return prng.choice(weighted_choices)
+
+
 def expand_decision(decision, prng):
     """ Take in a decision and a random number generator.  Expand variables in
         decision's values and headers until all values are fully expanded and