Add common IntegrationTest class with assertions and RunAllTests() func

tools/chrome_proxy/webdriver/common.py

 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import argparse
 import json
 import os
 import re
 import socket
 import shlex
 import sys
 import time
 import traceback
+import unittest
 
 sys.path.append(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir,
   os.pardir, 'third_party', 'webdriver', 'pylib'))
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
 
 # TODO(robertogden): Add logging.
 
 
 def ParseFlags():
   """Parses the given command line arguments.
 
   Returns:
     A new Namespace object with class properties for each argument added below.
     See pydoc for argparse.
   """
   parser = argparse.ArgumentParser()
-  parser.add_argument('--browser_args', nargs=1, type=str, help='Override '
-    'browser flags in code with these flags')
-  parser.add_argument('--via_header_value', metavar='via_header', nargs=1,
+  parser.add_argument('--browser_args', type=str, help='Override browser flags '
+    'in code with these flags')
+  parser.add_argument('--via_header_value',
     default='1.1 Chrome-Compression-Proxy', help='What the via should match to '
     'be considered valid')
   parser.add_argument('--android', help='If given, attempts to run the test on '
     'Android via adb. Ignores usage of --chrome_exec', action='store_true')
-  parser.add_argument('--android_package', nargs=1,
+  parser.add_argument('--android_package',
     default='com.android.chrome', help='Set the android package for Chrome')
-  parser.add_argument('--chrome_exec', nargs=1, type=str, help='The path to '
+  parser.add_argument('--chrome_exec', type=str, help='The path to '
     'the Chrome or Chromium executable')
-  parser.add_argument('chrome_driver', nargs=1, type=str, help='The path to '
+  parser.add_argument('chrome_driver', type=str, help='The path to '
     'the ChromeDriver executable. If not given, the default system chrome '
     'will be used.')
   parser.add_argument('--disable_buffer', help='Causes stdout and stderr from '
     'tests to output normally. Otherwise, the standard output and standard '
     'error streams are buffered during the test run, and output during a '
     'passing test is discarded. Output will always be echoed normally on test '
     'fail or error and is added to the failure messages.', action='store_true')
   parser.add_argument('-c', '--catch', help='Control-C during the test run '
     'waits for the current test to end and then reports all the results so '
     'far. A second Control-C raises the normal KeyboardInterrupt exception.',
     action='store_true')
   parser.add_argument('-f', '--failfast', help='Stop the test run on the first '
     'error or failure.', action='store_true')
   # TODO(robertogden): Log sys.argv here.
   return parser.parse_args(sys.argv[1:])
 
-def HandleException(test_name=None):
-  """Writes the exception being handled and a stack trace to stderr.
-
-  Args:
-    test_name: The string name of the test that led to this exception.
-  """
-  sys.stderr.write("**************************************\n")
-  sys.stderr.write("**************************************\n")
-  sys.stderr.write("**                                  **\n")
-  sys.stderr.write("**       UNCAUGHT EXCEPTION         **\n")
-  sys.stderr.write("**                                  **\n")
-  sys.stderr.write("**************************************\n")
-  sys.stderr.write("**************************************\n")
-  if test_name:
-    sys.stderr.write("Failed test: %s" % test_name)
-  traceback.print_exception(*sys.exc_info())
-  sys.exit(1)
 
 class TestDriver:
   """The main driver for an integration test.
 
   This class is the tool that is used by every integration test to interact with
   the Chromium browser and validate proper functionality. This class sits on top
   of the Selenium Chrome Webdriver with added utility and helper functions for
   Chrome-Proxy. This class should be used with Python's 'with' operator.
 
   Attributes:
@@ skipping 25 matching lines @@
     override the argument.
     """
     def GetDictKey(argument):
       return argument.split('=', 1)[0]
     if self._flags.browser_args and len(self._flags.browser_args) > 0:
       # Build a dict of flags mapped to the whole argument.
       original_args = {}
       for arg in self._chrome_args:
           original_args[GetDictKey(arg)] = arg
       # Override flags given in code with any command line arguments.
-      for override_arg in shlex.split(self._flags.browser_args[0]):
+      for override_arg in shlex.split(self._flags.browser_args):
         arg_key = GetDictKey(override_arg)
         if arg_key in original_args:
           self._chrome_args.remove(original_args[arg_key])
         self._chrome_args.add(override_arg)
     # Always add the flag that allows histograms to be queried in javascript.
     self._chrome_args.add('--enable-stats-collection-bindings')
 
   def _StartDriver(self):
     """Parses the flags to pass to Chromium, then starts the ChromeDriver.
 
     If running Android, the Android package name is passed to ChromeDriver here.
     """
     self._OverrideChromeArgs()
     capabilities = {
       'loggingPrefs': {'performance': 'INFO'},
       'chromeOptions': {
         'args': list(self._chrome_args)
       }
     }
     if self._flags.android:
       capabilities['chromeOptions'].update({
         'androidPackage': self._flags.android_package,
       })
     elif self._flags.chrome_exec:
       capabilities['chrome.binary'] = self._flags.chrome_exec
-    driver = webdriver.Chrome(executable_path=self._flags.chrome_driver[0],
+    driver = webdriver.Chrome(executable_path=self._flags.chrome_driver,
       desired_capabilities=capabilities)
     driver.command_executor._commands.update({
       'getAvailableLogTypes': ('GET', '/session/$sessionId/log/types'),
       'getLog': ('POST', '/session/$sessionId/log')})
     self._driver = driver
 
   def _StopDriver(self):
     """Nicely stops the ChromeDriver.
     """
     self._driver.quit()
@@ skipping 198 matching lines @@
     _url: the fetched url
     _protocol: The protocol used to get the response.
     _port: The remote port number used to get the response.
     _status: The integer status code of the response
     _request_type: What caused this request (Document, XHR, etc)
     _flags: A Namespace object from ParseFlags()
   """
 
   def __init__(self, response_headers, request_headers, url, protocol, port,
       status, request_type):
-    self._response_headers = response_headers
-    self._request_headers = request_headers
+    self._response_headers = {}
+    self._request_headers = {}
     self._url = url
     self._protocol = protocol
     self._port = port
     self._status = status
     self._request_type = request_type
     self._flags = ParseFlags()
+    # Make all header names lower case.
+    for name in response_headers:
+      self._response_headers[name.lower()] = response_headers[name]
+    for name in request_headers:
+      self._request_headers[name.lower()] = request_headers[name]
 
   def __str__(self):
     self_dict = {
       'response_headers': self._response_headers,
       'request_headers': self._request_headers,
       'url': self._url,
       'protocol': self._protocol,
       'port': self._port,
       'status': self._status,
       'request_type': self._request_type
@@ skipping 33 matching lines @@
       self._flags.via_header_value)
 
   def WasXHR(self):
     return self.request_type == 'XHR'
 
   def UsedHTTP(self):
     return self._protocol == 'http/1.1'
 
   def UsedHTTP2(self):
     return self._protocol == 'h2'
+
+class IntegrationTest(unittest.TestCase):
+  """This class adds ChromeProxy-specific assertions to the generic
+  unittest.TestCase class.
+  """
+
+  def assertHasChromeProxyViaHeader(self, http_response):
+    """Asserts that the Via header in the given HTTPResponse matches the
+    expected value as given on the command line.
+
+    Args:
+      http_response: The HTTPResponse object to check.
+    """
+    expected_via_header = ParseFlags().via_header_value
+    self.assertIn('via', http_response.response_headers)
+    self.assertEqual(expected_via_header, http_response.response_headers['via'])
+
+  def assertNotHasChromeProxyViaHeader(self, http_response):
+    """Asserts that the Via header in the given HTTPResponse does not match the
+    expected value as given on the command line.
+
+    Args:
+      http_response: The HTTPResponse object to check.
+    """
+    expected_via_header = ParseFlags().via_header_value
+    self.assertNotIn('via', http_response.response_headers)
+    if 'via' in http_response.response_headers:
+      self.assertNotIn(expected_via_header,
+        http_response.response_headers['via'])
+
+  @staticmethod
+  def RunAllTests():
+    """A simple helper method to run all tests using unittest.main().
+    """
+    # The unittest library uses sys.argv itself and is easily confused by our
+    # command line options. Pass it a simpler argv instead, while working in the
+    # unittest command line args functionality.
+    flags = ParseFlags()
+    unittest.main(argv=[sys.argv[0]], verbosity=2, failfast=flags.failfast,
+      catchbreak=flags.catch, buffer=(not flags.disable_buffer))