3 files changed, 71 insertions, 3 deletions

diff --git a/scripts/generator/cpp/ast.py b/scripts/generator/cpp/ast.py
index 6f61f877..bb8226d7 100755
--- a/scripts/generator/cpp/ast.py
+++ b/scripts/generator/cpp/ast.py
@@ -1546,7 +1546,7 @@ class AstBuilder(object):
             self._AddBackToken(token)
 
         return class_type(class_token.start, class_token.end, class_name,
-                          bases, None, body, self.namespace_stack)
+                          bases, templated_types, body, self.namespace_stack)
 
     def handle_namespace(self):
         token = self._GetNextToken()
diff --git a/scripts/generator/cpp/gmock_class.py b/scripts/generator/cpp/gmock_class.py
index 427d206a..3c8a877d 100755
--- a/scripts/generator/cpp/gmock_class.py
+++ b/scripts/generator/cpp/gmock_class.py
@@ -88,7 +88,11 @@ def _GenerateMethods(output_lines, source, class_node):
           if source[first_param.start:first_param.end].strip() == 'void':
             # We must treat T(void) as a function with no parameters.
             num_parameters = 0
-      mock_method_macro = 'MOCK_%sMETHOD%d' % (const, num_parameters)
+      tmpl = ''
+      if class_node.templated_types:
+        tmpl = '_T'
+      mock_method_macro = 'MOCK_%sMETHOD%d%s' % (const, num_parameters, tmpl)
+
       args = ''
       if node.parameters:
         # Due to the parser limitations, it is impossible to keep comments
@@ -126,6 +130,7 @@ def _GenerateMocks(filename, source, ast_list, desired_class_names):
         # desired_class_names being None means that all classes are selected.
         (not desired_class_names or node.name in desired_class_names)):
       class_name = node.name
+      parent_name = class_name
       processed_class_names.add(class_name)
       class_node = node
       # Add namespace before the class.
@@ -133,8 +138,21 @@ def _GenerateMocks(filename, source, ast_list, desired_class_names):
         lines.extend(['namespace %s {' % n for n in class_node.namespace])  # }
         lines.append('')
 
+      # Add template args for templated classes.
+      if class_node.templated_types:
+        # TODO(paulchang): The AST doesn't preserve template argument order,
+        # so we have to make up names here.
+        # TODO(paulchang): Handle non-type template arguments (e.g.
+        # template<typename T, int N>).
+        template_arg_count = len(class_node.templated_types.keys())
+        template_args = ['T%d' % n for n in range(template_arg_count)]
+        template_decls = ['typename ' + arg for arg in template_args]
+        lines.append('template <' + ', '.join(template_decls) + '>')
+        parent_name += '<' + ', '.join(template_args) + '>'
+
       # Add the class prolog.
-      lines.append('class Mock%s : public %s {' % (class_name, class_name))  # }
+      lines.append('class Mock%s : public %s {'  # }
+                   % (class_name, parent_name))
       lines.append('%spublic:' % (' ' * (_INDENT // 2)))
 
       # Add all the methods.
diff --git a/scripts/generator/cpp/gmock_class_test.py b/scripts/generator/cpp/gmock_class_test.py
index 7aa70276..07d59571 100755
--- a/scripts/generator/cpp/gmock_class_test.py
+++ b/scripts/generator/cpp/gmock_class_test.py
@@ -196,6 +196,18 @@ class Foo {
         'MOCK_METHOD0(Bar,\nmap<int, string>());',
         self.GenerateMethodSource(source))
 
+  def testSimpleMethodInTemplatedClass(self):
+    source = """
+template<class T>
+class Foo {
+ public:
+  virtual int Bar();
+};
+"""
+    self.assertEqualIgnoreLeadingWhitespace(
+        'MOCK_METHOD0_T(Bar,\nint());',
+        self.GenerateMethodSource(source))
+
 
 class GenerateMocksTest(TestCase):
 
@@ -255,5 +267,43 @@ void());
     self.assertEqualIgnoreLeadingWhitespace(
         expected, self.GenerateMocks(source))
 
+  def testTemplatedForwardDeclaration(self):
+    source = """
+template <class T> class Forward;  // Forward declaration should be ignored.
+class Test {
+ public:
+  virtual void Foo();
+};
+"""
+    expected = """\
+class MockTest : public Test {
+public:
+MOCK_METHOD0(Foo,
+void());
+};
+"""
+    self.assertEqualIgnoreLeadingWhitespace(
+        expected, self.GenerateMocks(source))
+
+  def testTemplatedClass(self):
+    source = """
+template <typename S, typename T>
+class Test {
+ public:
+  virtual void Foo();
+};
+"""
+    expected = """\
+template <typename T0, typename T1>
+class MockTest : public Test<T0, T1> {
+public:
+MOCK_METHOD0_T(Foo,
+void());
+};
+"""
+    self.assertEqualIgnoreLeadingWhitespace(
+        expected, self.GenerateMocks(source))
+
+
 if __name__ == '__main__':
   unittest.main()