Generalize rebase-errors.py so it works for other files.

net/data/verify_certificate_chain_unittest/rebase-errors.py

Index: net/data/verify_certificate_chain_unittest/rebase-errors.py
diff --git a/net/data/verify_certificate_chain_unittest/rebase-errors.py b/net/data/verify_certificate_chain_unittest/rebase-errors.py
index 0ac03b6600ac7f581ceb0a2f3057e0486ef2f64c..dfc9d93751d7ac199a6a541c084a5340aee82ba4 100755
--- a/net/data/verify_certificate_chain_unittest/rebase-errors.py
+++ b/net/data/verify_certificate_chain_unittest/rebase-errors.py
@@ -14,19 +14,47 @@ To use this run the affected tests, and then pass the input to this script
   $ ./out/Release/net_unittests --gtest_filter="*VerifyCertificateChain*" | \
      net/data/verify_certificate_chain_unittest/rebase-errors.py
 
-The script will search the unit-test (results.txt in above example) and look
-for failure lines and the corresponding actual error string.
+The script works by scanning the stdout looking for gtest failures when
+comparing "errors.ToDebugString()". The C++ test side should have been
+instrumented to dump out the test file's path on mismatch.
 
-It will then go and update the corresponding .pem and .py file.
+This script will then update the corresponding file(s) -- a .pem file, and
+possibly an accompanying .py file.
 """
 
 import common
-import glob
 import os
 import sys
 import re
 
 
+# Regular expression to find the failed errors in test stdout.
+#  * Group 1 of the match is the actual error text (backslash-escaped)
+#  * Group 2 of the match is file path (relative to //src) where the expected
+#    errors were read from.
+failed_test_regex = re.compile(r"""
+Value of: errors.ToDebugString\(\)
+  Actual: "(.*)"
+(?:.|\n)+?
+Test file: (.*)
+""", re.MULTILINE)
+
+
+# Regular expression to find the ERRORS block (and any text above it) in a PEM
+# file. The assumption is that ERRORS is not the very first block in the file
+# (since it looks for an -----END to precede it).
+#  * Group 1 of the match is the ERRORS block content and any comments
+#    immediately above it.
+errors_block_regex = re.compile(r""".*
+-----END .*?-----
+
+(.*?
+-----BEGIN ERRORS-----
+.*?
+-----END ERRORS-----
+)""", re.MULTILINE | re.DOTALL)
+
+
 def read_file_to_string(path):
   """Reads a file entirely to a string"""
   with open(path, 'r') as f:
@@ -40,36 +68,17 @@ def write_string_to_file(data, path):
     f.write(data)
 
 
-def get_file_paths_for_test(test_name):
-  """Returns the file paths (as a tuple) that define a particular unit test.
-  For instance given test name 'IntermediateLacksBasicConstraints' it would
-  return the paths to:
-
-   * intermediate-lacks-basic-constraints.pem,
-   * generate-intermediate-lacks-basic-constraints.py
-  """
-  # The directory that this python script is stored in.
-  base_dir = os.path.dirname(os.path.realpath(__file__))
+def get_py_path(pem_path):
+  """Returns the .py filepath used to generate the given .pem path, which may
+  or may not exist.
 
-  # The C++ test name is just a camel case verson of the file name. Rather than
-  # converting directly from camel case to a file name, it is simpler to just
-  # scan the file list and see which matches. (Not efficient but good enough).
-  paths = glob.glob(os.path.join(base_dir, '*.pem'))
-
-  for pem_path in paths:
-    file_name = os.path.basename(pem_path)
-    file_name_no_extension = os.path.splitext(file_name)[0]
-
-    # Strip the hyphens in file name to bring it closer to the camel case.
-    transformed = file_name_no_extension.replace('-', '')
-
-    # Now all that differs is the case.
-    if transformed.lower() == test_name.lower():
-      py_file_name = 'generate-' + file_name_no_extension + '.py'
-      py_path = os.path.join(base_dir, py_file_name)
-      return (pem_path, py_path)
-
-  return None
+  Some test files (notably those in verify_certificate_chain_unittest/ have a
+  "generate-XXX.py" script that builds the "XXX.pem" file. Build the path to
+  the corresponding "generate-XXX.py" (which may or may not exist)."""
+  file_name = os.path.basename(pem_path)
+  file_name_no_extension = os.path.splitext(file_name)[0]
+  py_file_name = 'generate-' + file_name_no_extension + '.py'
+  return os.path.join(os.path.dirname(pem_path), py_file_name)
 
 
 def replace_string(original, start, end, replacement):
@@ -81,14 +90,14 @@ def fixup_pem_file(path, actual_errors):
   """Updates the ERRORS block in the test .pem file"""
   contents = read_file_to_string(path)
 
-  # This assumes that ERRORS is the last thing in file, and comes after the
-  # VERIFY_RESULT block.
-  kEndVerifyResult = '-----END VERIFY_RESULT-----'
-  contents = contents[0:contents.index(kEndVerifyResult)]
-  contents += kEndVerifyResult
-  contents += '\n'
-  contents += '\n'
-  contents += common.text_data_to_pem('ERRORS', actual_errors)
+  m = errors_block_regex.search(contents)
+
+  if not m:
+    print "Couldn't find ERRORS block in %s" % (path)
+    return
+
+  contents = replace_string(contents, m.start(1), m.end(1),
+                            common.text_data_to_pem('ERRORS', actual_errors))
 
   # Update the file.
   write_string_to_file(contents, path)
@@ -99,7 +108,7 @@ def fixup_py_file(path, actual_errors):
   contents = read_file_to_string(path)
 
   # This assumes that the errors variable uses triple quotes.
-  prog = re.compile(r'^errors = """(.*)"""', re.MULTILINE | re.DOTALL)
+  prog = re.compile(r'^errors = """(.*?)"""', re.MULTILINE | re.DOTALL)
   result = prog.search(contents)
 
   # Replace the stuff in between the triple quotes with the actual errors.
@@ -110,21 +119,39 @@ def fixup_py_file(path, actual_errors):
   write_string_to_file(contents, path)
 
 
-def fixup_test(test_name, actual_errors):
-  """Updates the test files used by |test_name|, setting the expected error to
-  |actual_errors|"""
+def get_src_root():
+  """Returns the path to the enclosing //src directory. This assumes the
+  current script is inside the source tree."""
+  cur_dir = os.path.dirname(os.path.realpath(__file__))
 
-  # Determine the paths for the corresponding *.pem file and generate-*.py
-  pem_path, py_path = get_file_paths_for_test(test_name)
+  while True:
+    parent_dir, dirname = os.path.split(cur_dir)
+    # Check if it looks like the src/ root.
+    if dirname == "src" and os.path.isdir(os.path.join(cur_dir, "net")):
+      return cur_dir
+    if not parent_dir or parent_dir == cur_dir:
+      break
+    cur_dir = parent_dir
 
-  fixup_pem_file(pem_path, actual_errors)
-  fixup_py_file(py_path, actual_errors)
+  print "Couldn't find src dir"
+  sys.exit(1)
+
+
+def get_abs_path(rel_path):
+  """Converts |rel_path| (relative to src) to a full path"""
+  return os.path.join(get_src_root(), rel_path)
+
+
+def fixup_errors_for_file(actual_errors, pem_path):
+  """Updates the errors in |test_file_path| (.pem file) to match
+  |actual_errors|"""
 
+  fixup_pem_file(pem_path, actual_errors)
 
-kTestNamePattern = (r'^\[ RUN      \] VerifyCertificateChain/'
-                    'VerifyCertificateChainSingleRootTest/0\.(.*)$')
-kValueOfLine = 'Value of: errors.ToDebugString()'
-kActualPattern = '^  Actual: "(.*)"$'
+  # If the test has a generator script update it too.
+  py_path = get_py_path(pem_path)
+  if os.path.isfile(py_path):
+    fixup_py_file(py_path, actual_errors)
 
 
 def main():
@@ -140,28 +167,11 @@ def main():
     print 'Reading input from stdin...'
     test_stdout = sys.stdin.read()
 
-  lines = test_stdout.split('\n')
-
-  # Iterate over each line of the unit test stdout.
-  for i in range(len(lines) - 3):
-    # Figure out the name of the test.
-    m = re.search(kTestNamePattern, lines[i])
-    if not m:
-      continue
-    test_name = m.group(1)
-
-    # Confirm that it is a failure having to do with the errors.
-    if lines[i + 2] != kValueOfLine:
-      continue
-
-    # Get the actual error text (which in gtest output is escaped).
-    m = re.search(kActualPattern, lines[i + 3])
-    if not m:
-      continue
-    actual = m.group(1)
-    actual = actual.decode('string-escape')
-
-    fixup_test(test_name, actual)
+  for m in failed_test_regex.finditer(test_stdout):
+    actual_errors = m.group(1)
+    actual_errors = actual_errors.decode('string-escape')
+    relative_test_path = m.group(2)
+    fixup_errors_for_file(actual_errors, get_abs_path(relative_test_path))
 
 
 if __name__ == "__main__":