1#!/usr/bin/env python
2# Copyright 2019 The Chromium Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""
7Unit tests for extractor.py.
8"""
9
10from __future__ import print_function
11
12import argparse
13import glob
14import os
15import unittest
16import re
17import subprocess
18
19
20def run_extractor(file, *extra_args):
21  script_path = os.path.join('..', 'extractor.py')
22  cmd_line = ("python", script_path, '--no-filter', file) + extra_args
23  return subprocess.Popen(
24      cmd_line, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
25
26
27def get_expected_files(source_file):
28  stdout_file = re.sub(r'\.cc$', '-stdout.txt', source_file)
29  stderr_file = re.sub(r'\.cc$', '-stderr.txt', source_file)
30  return (stdout_file, stderr_file)
31
32
33def dos2unix(str):
34  """Convers CRLF to LF."""
35  return str.replace('\r\n', '\n')
36
37
38def remove_tracebacks(str):
39  """Removes python tracebacks from the string."""
40  regex = re.compile(
41      r'''
42       # A line that says "Traceback (...):"
43       ^Traceback[^\n]*:\n
44       # Followed by lines that begin with whitespace.
45       ((\s.*)?\n)*''',
46      re.MULTILINE | re.VERBOSE)
47  return re.sub(regex, '', str)
48
49
50class ExtractorTest(unittest.TestCase):
51  def testExtractor(self):
52    for source_file in glob.glob('*.cc'):
53      print("Running test on %s..." % source_file)
54      (stdout_file, stderr_file) = get_expected_files(source_file)
55      with open(stdout_file) as f:
56        expected_stdout = dos2unix(f.read())
57      with open(stderr_file) as f:
58        expected_stderr = dos2unix(f.read())
59
60      proc = run_extractor(source_file)
61      (stdout, stderr) = map(dos2unix, proc.communicate())
62
63      self.assertEqual(expected_stderr, remove_tracebacks(stderr))
64      expected_returncode = 2 if expected_stderr else 0
65      self.assertEqual(expected_returncode, proc.returncode)
66      self.assertEqual(expected_stdout, stdout)
67
68
69def generate_expected_files():
70  for source_file in glob.glob('*.cc'):
71    proc = run_extractor(source_file)
72    (stdout, stderr) = proc.communicate()
73
74    (stdout_file, stderr_file) = get_expected_files(source_file)
75    with open(stdout_file, "w") as f:
76      f.write(stdout)
77    with open(stderr_file, "w") as f:
78      f.write(remove_tracebacks(stderr))
79
80
81if __name__ == '__main__':
82  parser = argparse.ArgumentParser(
83      description="Network Traffic Annotation Extractor Unit Tests")
84  parser.add_argument(
85      '--generate-expected-files', action='store_true',
86      help='Generate "-stdout.txt" and "-stderr.txt" for file in test_data')
87  args = parser.parse_args()
88
89  # Set directory for both test and gen command to the test_data folder.
90  os.chdir(os.path.join(os.path.dirname(__file__), 'test_data'))
91
92  # Run the extractor script with --generate-compdb to ensure the
93  # compile_commands.json file exists in the default output directory.
94  proc = run_extractor(os.devnull, '--generate-compdb')
95  proc.communicate()  # Wait until extractor finishes running.
96
97  if args.generate_expected_files:
98    generate_expected_files()
99  else:
100    unittest.main()
101