Allow parametrize argvalues and ids to accept generators

_pytest/python.py

@@ -837,7 +837,7 @@ def parametrize(self, argnames, argvalues, indirect=False, ids=None,
         :arg argnames: a comma-separated string denoting one or more argument
                        names, or a list/tuple of argument strings.
 
-        :arg argvalues: The list of argvalues determines how often a
+        :arg argvalues: The list/generator of argvalues determines how often a
             test is invoked with different argument values.  If only one
             argname was specified argvalues is a list of simple values.  If N
             argnames were specified, argvalues must be a list of N-tuples,
@@ -849,7 +849,7 @@ def parametrize(self, argnames, argvalues, indirect=False, ids=None,
             function so that it can perform more expensive setups during the
             setup phase of a test rather than at collection time.
 
-        :arg ids: list of string ids, or a callable.
+        :arg ids: list of string ids, generator or a callable.
             If strings, each is corresponding to the argvalues so that they are
             part of the test id.
             If callable, it should take one argument (a single argvalue) and return
@@ -869,6 +869,10 @@ def parametrize(self, argnames, argvalues, indirect=False, ids=None,
         # at Function init
         newkeywords = {}
         unwrapped_argvalues = []
+
+        if argvalues and not hasattr(argvalues, '__len__'):
+            argvalues = list(argvalues)
+
         for i, argval in enumerate(argvalues):
             while isinstance(argval, MarkDecorator):
                 newmark = MarkDecorator(argval.markname,
@@ -900,6 +904,8 @@ def parametrize(self, argnames, argvalues, indirect=False, ids=None,
         if callable(ids):
             idfn = ids
             ids = None
+        if ids and not hasattr(ids, '__len__'):
+            ids = list(ids)
         if ids and len(ids) != len(argvalues):
             raise ValueError('%d tests specified with %d ids' %(
                              len(argvalues), len(ids)))

testing/python/metafunc.py

@@ -342,6 +342,50 @@ def test_3(self, arg, arg2):
             *6 passed*
         """)
 
+    def test_parametrize_generators(self):
+        def generator_ids():
+            for value in ['1', '2']:
+                yield value
+
+        def generator_argvalues():
+            for value in [(1, 2), (3, 4)]:
+                yield value
+
+        metafunc = self.Metafunc(lambda x, y: None)
+        metafunc.parametrize(("x", "y"), generator_argvalues(), ids=generator_ids())
+        for i in [1, 2]:
+            ids = [x.id for x in metafunc._calls]
+            assert ids == ['1', '2']
+            assert len(metafunc._calls) == 2
+            assert metafunc._calls[0].funcargs == dict(x=1, y=2)
+            assert metafunc._calls[1].funcargs == dict(x=3, y=4)
+
+    def test_parametrize_generators_multiple_times(self, testdir):
+        testdir.makepyfile("""
+            import pytest
+
+            def generator_ids():
+                for value in ['1']:
+                    yield value
+
+            def generator_argvalues():
+                for value in [(1)]:
+                    yield value
+
+
+            @pytest.mark.parametrize("x", generator_argvalues(), ids=generator_ids())
+            def test_func(x):
+                assert 1 == x
+
+            class TestClass:
+                @pytest.mark.parametrize("y", generator_argvalues(), ids=generator_ids())
+                def test_meth(self, y):
+                    assert 1 == y
+        """)
+        result = testdir.runpytest()
+        assert result.ret == 0
+        result.assert_outcomes(passed=2)
+
 class TestMetafuncFunctional:
     def test_attributes(self, testdir):
         p = testdir.makepyfile("""