Fix run_threaded decorator so that it works with named_parameters

alimuldal · alimuldal · commit 96db220cd0c6 · 2018-01-26T14:14:17.000Z
PiperOrigin-RevId: 182068715
diff --git a/dm_control/mujoco/testing/decorators.py b/dm_control/mujoco/testing/decorators.py
@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import sys
 import threading
 
@@ -42,6 +43,7 @@ def run_threaded(num_threads=4, calls_per_thread=10):
   """
   def decorator(test_method):
     """Decorator around the test method."""
+    @functools.wraps(test_method)  # Needed for `named_parameters` to work.
     def decorated_method(self, *args, **kwargs):
       """Actual method this factory will return."""
       exceptions = []
diff --git a/dm_control/mujoco/testing/decorators_test.py b/dm_control/mujoco/testing/decorators_test.py
@@ -22,6 +22,7 @@
 # Internal dependencies.
 
 from absl.testing import absltest
+from absl.testing import parameterized
 
 from dm_control.mujoco.testing import decorators
 import mock
@@ -42,7 +43,9 @@ def test_number_of_threads(self, mock_threading):
     mock_threading.Thread = mock.MagicMock(side_effect=mock_threads)
 
     test_decorator = decorators.run_threaded(num_threads=num_threads)
-    test_runner = test_decorator(mock.MagicMock())
+    tested_method = mock.MagicMock()
+    tested_method.__name__ = "foo"
+    test_runner = test_decorator(tested_method)
     test_runner(self)
 
     for thread in mock_threads:
@@ -53,13 +56,39 @@ def test_number_of_iterations(self):
     calls_per_thread = 5
 
     tested_method = mock.MagicMock()
+    tested_method.__name__ = "foo"
     test_decorator = decorators.run_threaded(
         num_threads=1, calls_per_thread=calls_per_thread)
     test_runner = test_decorator(tested_method)
     test_runner(self)
 
     self.assertEqual(calls_per_thread, tested_method.call_count)
 
+  def test_works_with_named_parameters(self):
+
+    func = mock.MagicMock()
+    names = ["foo", "bar", "baz"]
+    params = [1, 2, 3]
+    calls_per_thread = 2
+    num_threads = 4
+
+    class FakeTest(parameterized.TestCase):
+
+      @parameterized.named_parameters(zip(names, params))
+      @decorators.run_threaded(calls_per_thread=calls_per_thread,
+                               num_threads=num_threads)
+      def test_method(self, param):
+        func(param)
+
+    suite = absltest.TestLoader().loadTestsFromTestCase(FakeTest)
+    suite.debug()  # Run tests without collecting the output.
+
+    expected_call_count = len(params) * calls_per_thread * num_threads
+
+    self.assertEqual(func.call_count, expected_call_count)
+    actual_params = {call[0][0] for call in func.call_args_list}
+    self.assertSetEqual(set(params), actual_params)
+
 
 if __name__ == "__main__":
   absltest.main()
