testing generator

Author: AaronRobson

tests/test_fizzbuzz.py

@@ -1,6 +1,7 @@
 #!/usr/bin/python
 
 import unittest
+from itertools import islice
 
 import fizzbuzz
 
@@ -17,5 +18,24 @@ def test(self):
         self.assertEqual(f(15), 'FizzBuzz')
 
 
+class TestFizzBuzzLoop(unittest.TestCase):
+
+    def test(self):
+        expected = [
+            '1',
+            '2',
+            'Fizz',
+            '4',
+            'Buzz',
+        ]
+        actual = take(len(expected), fizzbuzz.fizzBuzzLoop())
+        self.assertEqual(expected, actual)
+
+
+def take(n, iterable):
+    "Return first n items of the iterable as a list"
+    return list(islice(iterable, n))
+
+
 if __name__ == "__main__":
     unittest.main()