python-tuf/tests/utils.py

#!/usr/bin/env python

# Copyright 2020, TUF contributors
# SPDX-License-Identifier: MIT OR Apache-2.0

"""
<Program Name>
  utils.py

<Started>
  August 3, 2020.

<Author>
  Jussi Kukkonen

<Copyright>
  See LICENSE-MIT OR LICENSE for licensing information.

<Purpose>
  Provide common utilities for TUF tests
"""

import argparse
import errno
import logging
import socket
import time
import subprocess
import tempfile
import random

import tuf.log

logger = logging.getLogger(__name__)

try:
  # is defined in Python 3
  TimeoutError
except NameError:
  # Define for Python 2
  class TimeoutError(Exception):

    def __init__(self, value="Timeout"):
      self.value = value

    def __str__(self):
      return repr(self.value)

# Wait until host:port accepts connections.
# Raises TimeoutError if this does not happen within timeout seconds
# There are major differences between operating systems on how this works
# but the current blocking connect() seems to work fast on Linux and seems
# to at least work on Windows (ECONNREFUSED unfortunately has a 2 second
# timeout on Windows)
def wait_for_server(host, port, timeout=10):
  start = time.time()
  remaining_timeout = timeout
  succeeded = False
  while not succeeded and remaining_timeout > 0:
    try:
      sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
      sock.settimeout(remaining_timeout)
      sock.connect((host, port))
      succeeded = True
    except socket.timeout as e:
      pass
    except IOError as e:
      # ECONNREFUSED is expected while the server is not started
      if e.errno not in [errno.ECONNREFUSED]:
        logger.warning("Unexpected error while waiting for server: " + str(e))
      # Avoid pegging a core just for this
      time.sleep(0.01)
    finally:
      if sock:
        sock.close()
        sock = None
      remaining_timeout = timeout - (time.time() - start)

  if not succeeded:
    raise TimeoutError


def configure_test_logging(argv):
  # parse arguments but only handle '-v': argv may contain 
  # other things meant for unittest argument parser
  parser = argparse.ArgumentParser(add_help=False)
  parser.add_argument('-v', '--verbose', action='count', default=0)
  args, _ = parser.parse_known_args(argv)
  
  if args.verbose <= 1:
    # 0 and 1 both mean ERROR: this way '-v' makes unittest print test
    # names without increasing log level
    loglevel = logging.ERROR
  elif args.verbose == 2:
    loglevel = logging.WARNING
  elif args.verbose == 3:
    loglevel = logging.INFO
  else:
    loglevel = logging.DEBUG

  logging.basicConfig(level=loglevel)
  tuf.log.set_log_level(loglevel)


class TestServerProcess():
  """
  <Purpose>
    Creates a child process with the subprocess.Popen object and
    TempFile object used for logging.

   <Arguments>
      log:
        Logger which will be used for logging.

      server:
        Path to the server to run in the subprocess.
        Default is "simpler_server.py".

      port:
        The port used to access the server. If none is provided,
        then one will be generated.
        Default is None.

      timeout:
        Time in seconds in which the server should start or otherwise
        TimeoutError error will be raised.
        If 0 is given, no check if the server has started will be done.
        Default is 10.

      popen_cwd:
        Current working directory used when instancing a
        subprocess.Popen object.
        Default is "."

      extra_cmd_args:
        List of additional arguments for the command
        which will start the subprocess.
        More precisely "python -u <path_to_server> <port> <extra_cmd_args>".
        Default is empty list.
  """


  def __init__(self, log, server='simple_server.py',
      port=None, timeout=10, popen_cwd=".",
      extra_cmd_args=[]):

    # Create temporary log file used for logging stdout and stderr
    # of the subprocess. In the mode "r+"" stands for reading and writing
    # and "t" stands for text mode.
    self.__temp_log_file = tempfile.TemporaryFile(mode='r+t')

    self.server = server
    self.port = port or random.randint(30000, 45000)
    self.__logger = log

    # The "-u" option forces stdin, stdout and stderr to be unbuffered.
    command = ['python', '-u', server, str(self.port)] + extra_cmd_args

    # We are reusing one server subprocess in multiple unit tests, but we are
    # collecting the logs per test.
    self.__server_process = subprocess.Popen(command,
        stdout=self.__temp_log_file, stderr=subprocess.STDOUT, cwd=popen_cwd)

    self.__logger.info('Server process with process id ' \
        + str(self.__server_process.pid) + " serving on port " \
        + str(self.port) + ' started.')

    if timeout > 0:
      try:
        wait_for_server('localhost', self.port, timeout)
      except Exception as e:
        self.__logger.error(
            "Error while waiting for the server to start: " + repr(e))
        # Make sure that errors from the server side will be logged.
        self.flush_log()
        raise e



  def flush_log(self):
    """Logs contents from TempFile, truncates buffer"""

    # Seek is needed to move the pointer to the beginning of the file, because
    # the subprocess could have read and/or write and thus moved the pointer.
    self.__temp_log_file.seek(0)
    log_message = self.__temp_log_file.read()

    if len(log_message) > 0:
      title = "Test server (" + self.server + ") output:"
      message = [title] + log_message.splitlines()
      self.__logger.info('\n| '.join(message))

      # Make sure the file is empty before the next test logs new information.
      self.__temp_log_file.truncate(0)



  def clean(self):
    """Kills the subprocess and closes the TempFile.
    Calls flush_log to check for logged information, but not yet flushed."""

    # If there is anything logged, flush it before closing the resourses.
    self.flush_log()

    self.__temp_log_file.close()

    if self.__server_process.returncode is None:
      self.__logger.info('Server process ' + str(self.__server_process.pid) +
          ' terminated.')
      self.__server_process.kill()
      self.__server_process.wait()