#!/usr/bin/env vpython3
# Copyright 2019 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""
Unit tests for extractor.py.
"""

from __future__ import print_function

import argparse
import os
import sys
import unittest
import re
import subprocess

from glob import glob
from pathlib import Path
from typing import Tuple

# Path to the directory where this script is.
SCRIPT_DIR = Path(__file__).resolve().parent


def run_extractor(file_path: Path) -> Tuple[bytes, bytes]:
  script_path = Path('../extractor.py')
  cmd_line = ('python3', str(script_path), '--no-filter', file_path)
  return subprocess.Popen(
    cmd_line, stdout=subprocess.PIPE, stderr=subprocess.PIPE
  )


def get_expected_files(source_file: Path) -> Path:
  stdout_file = source_file.with_name(
    '%s-stdout' % source_file.stem
  ).with_suffix('.txt')
  stderr_file = source_file.with_name(
    '%s-stderr' % source_file.stem
  ).with_suffix('.txt')
  return (stdout_file, stderr_file)


def remove_tracebacks(body: str):
  """Removes python tracebacks from the string."""
  regex = re.compile(
    r'''
       # A line that says "Traceback (...):"
       ^Traceback[^\n]*:\n
       # Followed by lines that begin with whitespace.
       ((\s.*)?\n)*''',
    re.MULTILINE | re.VERBOSE,
  )
  return re.sub(regex, '', body)


class ExtractorTest(unittest.TestCase):
  def testExtractor(self):
    files = list(Path(f) for f in glob('*.cc') + glob('*.java') + glob('*.h'))
    for source_file in files:
      if source_file.stem.startswith('test_'):
        continue

      print("Running test on %s..." % source_file)
      (stdout_file, stderr_file) = get_expected_files(source_file)
      expected_stdout = stdout_file.read_text()
      expected_stderr = stderr_file.read_text()

      proc = run_extractor(source_file)
      (stdout, stderr) = (b.decode() for b in proc.communicate())

      self.assertEqual(expected_stderr, remove_tracebacks(stderr))
      expected_returncode = 2 if expected_stderr else 0
      self.assertEqual(expected_returncode, proc.returncode)
      self.assertEqual(expected_stdout, stdout)


def generate_expected_files():
  files = list(Path(f) for f in glob('*.cc') + glob('*.java'))
  for source_file in files:
    if source_file.stem.startswith('test_'):
      continue
    proc = run_extractor(source_file)
    (stdout, stderr) = (b.decode() for b in proc.communicate())

    (stdout_file, stderr_file) = get_expected_files(source_file)
    stdout_file.write_text(stdout, newline='\n')
    stderr_file.write_text(stderr, newline='\n')


if __name__ == '__main__':
  parser = argparse.ArgumentParser(
    description="Network Traffic Annotation Extractor Unit Tests"
  )
  parser.add_argument(
    '--generate-expected-files',
    action='store_true',
    help='Generate "-stdout.txt" and "-stderr.txt" for file in test_data',
  )
  args, unknown_args = parser.parse_known_args()

  # Set directory for both test and gen command to the test_data folder.
  os.chdir(Path(__file__).resolve().parent / 'test_data')

  # XXX
  # # Run the extractor script with --generate-compdb to ensure the
  # # compile_commands.json file exists in the default output directory.
  # proc = run_extractor(os.devnull, '--generate-compdb')
  # proc.communicate()  # Wait until extractor finishes running.

  if args.generate_expected_files:
    generate_expected_files()
  else:
    unittest.main()