1 files changed, 56 insertions, 1 deletions

diff --git a/net/test/util.py b/net/test/util.py
index bed3e1d..cbcd2d0 100644
--- a/net/test/util.py
+++ b/net/test/util.py
@@ -13,4 +13,59 @@
 # limitations under the License.
 
 def GetPadLength(block_size, length):
-  return (block_size - (length % block_size)) % block_size
\ No newline at end of file
+  return (block_size - (length % block_size)) % block_size
+
+
+def InjectParameterizedTest(cls, param_list, name_generator):
+  """Injects parameterized tests into the provided class
+
+  This method searches for all tests that start with the name "ParamTest",
+  and injects a test method for each set of parameters in param_list. Names
+  are generated via the use of the name_generator.
+
+  Args:
+    cls: the class for which to inject all parameterized tests
+    param_list: a list of tuples, where each tuple is a combination of
+        of parameters to test (i.e. representing a single test case)
+    name_generator: A function that takes a combination of parameters and
+        returns a string that identifies the test case.
+  """
+  param_test_names = [name for name in dir(cls) if name.startswith("ParamTest")]
+
+  # Force param_list to an actual list; otherwise itertools.Product will hit
+  # the end, resulting in only the first ParamTest* method actually being
+  # parameterized
+  param_list = list(param_list)
+
+  # Parameterize each test method starting with "ParamTest"
+  for test_name in param_test_names:
+    func = getattr(cls, test_name)
+
+    for params in param_list:
+      # Give the test method a readable, debuggable name.
+      param_string = name_generator(*params)
+      new_name = "%s_%s" % (func.__name__.replace("ParamTest", "test"),
+                            param_string)
+      new_name = new_name.replace("(", "-").replace(")", "")  # remove parens
+
+      # Inject the test method
+      setattr(cls, new_name, _GetTestClosure(func, params))
+
+
+def _GetTestClosure(func, params):
+  """ Creates a no-argument test method for the given function and parameters.
+
+  This is required to be separate from the InjectParameterizedTest method, due
+  to some interesting scoping issues with internal function declarations. If
+  left in InjectParameterizedTest, all the tests end up using the same
+  instance of TestClosure
+
+  Args:
+    func: the function for which this test closure should run
+    params: the parameters for the run of this test function
+  """
+
+  def TestClosure(self):
+    func(self, *params)
+
+  return TestClosure