Commits

Anonymous committed 26a4306

Added test for eval_branches

Comments (0)

Files changed (2)

 
 from pgen import *
 
+
+class TestUtilFunctions(unittest.TestCase):
+    def testEvalBranches(self):
+        class Random(object):
+            def random(self):
+                return 0.9
+
+        rng = Random()
+
+        result = eval_branches(rng, [(0.5, 0), (1.0, 1)])
+        self.assertEqual(result, 1)
     
 class TestFunctionGenerator(unittest.TestCase):
 
 from pygen.cgen import *
 
 def eval_branches(rng, branches):
-    total = sum((x[0] for x in branches))
+    total = sum((chance for chance,_ in branches))
     val = rng.random() * total
 
     for chance, result in branches:
             return result
         else:
             val -= chance
-    return None
 
 class FunctionGenerator(object):
     def generate_arguments(self, args_num):