# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Stress test utility for repeating actions repeatedly on android devices.

Configures multiple devices to simultaneously run through the same set of
actions over and over, while keeping logs from various sources. Primarily
designed for playing audio to the devices and scanning their log output for
events, while running other adb commands in between.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime
from email import encoders
from email.mime import text
import email.mime.base as base
import email.mime.multipart as multipart
import logging
import mimetypes
import os
import platform
import re
import shlex
import signal
import smtplib
import socket
import subprocess
import sys
import tempfile
import threading
import time
import uuid
import wave
from absl import app
from absl import flags
import pexpect
import queue
import stress_test_common
import stress_test_pb2
from google.protobuf import text_format

_SUMMARY_LINES = "-" * 73

if sys.platform.startswith("win"):
  pexpect = None

_SUMMARY_COLUMNS = (
    "|        Event Type       |      Event Count     | Consecutive no event |")
_SUMMARY_COL_FORMATT = "|%-25.25s|% 22d|% 22d|"

FLAGS = flags.FLAGS
flags.DEFINE_string("notification_address", "",
                    "Email address where to send notification events. Will "
                    "default to $USER@google.com if not provided. No emails "
                    "will be sent if suppress_notification_emails is True.")
flags.DEFINE_bool("suppress_notification_emails", False,
                  "Prevents emails from being sent as notifications if True.")
flags.DEFINE_string("test_name", None,
                    "Name of stress test to run. For example, if you set this "
                    "to 'dsp_trigger_sw_rejection', the stress test in "
                    "'stress_test.dsp_trigger_sw_rejection.ascii_proto' will "
                    "be loaded and executed.")
# flags.mark_flag_as_required("test_name")
flags.DEFINE_string("output_root", "./",
                    "Path where directory should be generated containing all "
                    "logs from devices and moved files.")
flags.DEFINE_integer("num_iterations", None,
                     "If set to a positive number, the number of iterations of "
                     "the stress test to run. Otherwise, the test runs "
                     "forever.")
flags.DEFINE_list("devices", [],
                  "Serial numbers of devices that should be included in the "
                  "stress test. If empty, all devices will be used.")
flags.DEFINE_integer("print_summary_every_n", 10,
                     "Prints the summary to the log file every n iterations.")

flags.DEFINE_string("email_sender_address", "",
                    "Account to use for sending notification emails.")
flags.DEFINE_string("email_sender_password", "",
                    "Password to use for notification email account.")
flags.DEFINE_string("email_smtp_server", "smtp.gmail.com",
                    "SMTP server to use for sending notification emails.")
flags.DEFINE_integer("email_smtp_port", 465,
                     "Port to use for the notification SMTP server.")
flags.DEFINE_integer("device_settle_time", 5,
                     "Time to wait for devices to settle.")
flags.DEFINE_bool("use_sox", platform.system() != "Windows",
                  "Use sox for playback, otherwise, attempt to use platform "
                  "specific features.")
flags.DEFINE_bool("attach_bugreport", True,
                  "Attach bugreport to email if test failed.")
flags.DEFINE_bool("delete_data_dir", False,
                  "If true, code will delete all the files generated by this "
                  "test at the end.")

if platform.system().startswith("CYGWIN"):
  FLAGS.device_settle_time = 30


def QueueWorker(worker_queue):
  while True:
    work = worker_queue.get()
    try:
      work()
    except:  # pylint:disable=bare-except
      logging.exception("Exception in worker queue - task remains uncompleted.")
    worker_queue.task_done()


def SendNotificationEmail(subject, body, bugreport=None):
  """Sends an email with the specified subject and body.

     Also attach bugreport if bugreport location is provided as argument

  Args:
    subject: Subject of the email.
    body: Body of the email.
    bugreport: If provided, it will be attach to the email.
  """
  if FLAGS.suppress_notification_emails:
    logging.info("Email with subject '%s' has been suppressed", subject)
    return
  try:
    # Assemble the message to send.
    recpient_address = FLAGS.notification_address
    message = multipart.MIMEMultipart("alternative")
    message["From"] = "Stress Test on %s" % socket.gethostname()
    message["To"] = recpient_address
    message["Subject"] = subject
    message.attach(text.MIMEText(body, "plain"))
    message.attach(text.MIMEText("<pre>%s</pre>" % body, "html"))

    if FLAGS.attach_bugreport and bugreport:
      # buildozer: disable=unused-variable
      ctype, _ = mimetypes.guess_type(bugreport)
      maintype, subtype = ctype.split("/", 1)
      with open(bugreport, "rb") as fp:
        att = base.MIMEBase(maintype, subtype)
        att.set_payload(fp.read())
        encoders.encode_base64(att)
        att.add_header("Content-Disposition", "attachment", filename=bugreport)
        message.attach(att)

    # Send the message from our special account.
    server = smtplib.SMTP_SSL(FLAGS.email_smtp_server, FLAGS.email_smtp_port)
    server.login(FLAGS.email_sender_address, FLAGS.email_sender_password)
    server.sendmail(FLAGS.email_sender_address, recpient_address,
                    message.as_string())
    server.quit()
    logging.info("Email with subject '%s' has been sent", subject)
  except:  # pylint:disable=bare-except
    logging.exception("Failed to send notification email")


class ProcessLogger(threading.Thread):

  class EventScanner(object):

    def __init__(self, name, process_name, regexes):
      """Struct to store the data about an event.

      Args:
        name: Name of event.
        process_name: Name of the process that is being logged.
        regexes: An iteratable of regex strings that indicate an event has
            happened.
      """

      self.name = name
      self.process_name = process_name
      self.searches = [re.compile(regex).search for regex in regexes]
      self.count = 0

    def ScanForEvent(self, line, lock=None):
      """Checks the line for matches. If found, updates the internal counter."""

      for search in self.searches:
        if search(line.decode("utf-8")):
          # Grab the lock (if provided), update the counter, and release it.
          if lock: lock.acquire()
          self.count += 1
          if lock: lock.release()
          logging.info("Event '%s' detected on %s", self.name,
                       self.process_name)

  def __init__(self, name, command, output, events,
               restart_process, repeats_output_when_opened):
    """Threaded class that monitors processes for events, and logs output.

    Args:
      name: The name of the process being logged.
      command: A list of arguments to be passed to the subprocess to execute.
      output: Name of output file to write process stdout to. If blank or None,
          will not be generated.
      events: An iterable of LoggingEventConfigs to look for in the output.
      restart_process: Restart the process if it terminates by itself. This
          should typically be true, but false for processes that only should be
          run once and have their output logged.
      repeats_output_when_opened: Set to true if the process will repeat the
          output of a previous call when it is restarted. This will prevent
          duplicate lines from being logged.
    """
    super(ProcessLogger, self).__init__()
    self.name = name
    self.command = command
    self.restart_process = restart_process
    self.repeats_output_when_opened = repeats_output_when_opened
    self.process = None
    self.lock = threading.Lock()
    self.looking = False

    # Compile the list of regexes that we're supposed to be looking for.
    self.events = []
    for event in events:
      self.events.append(ProcessLogger.EventScanner(event.name, self.name,
                                                    event.regex))

    if output:
      stress_test_common.MakeDirsIfNeeded(os.path.dirname(output))
      self.output_fp = open(output, "w", encoding="utf-8")
      logging.info("Logging device info to %s", output)
    else:
      self.output_fp = None

  def GetEventCountsSinceLastCall(self):
    """Returns the counts of all events since this method was last called."""
    event_map = {}
    self.lock.acquire()
    for event in self.events:
      event_map[event.name] = event.count
      event.count = 0
    self.lock.release()
    return event_map

  def run(self):
    last_line = None
    should_log = True
    first_run = True
    skip_exception_line = False
    self.lock.acquire()
    last_run_time = 0
    while self.restart_process:
      self.lock.release()
      if not first_run:
        logging.info("Restarting process %s", "".join(str(self.command)))
        time_since_last_run = datetime.datetime.now() - last_run_time
        if time_since_last_run.total_seconds() < 1.0:
          needed_delay = 1.0 - time_since_last_run.total_seconds()
          logging.info("Delaying for %.2f seconds", needed_delay)
          time.sleep(needed_delay)
      else:
        first_run = False

      try:
        if pexpect:
          self.process = pexpect.spawn(" ".join(self.command), timeout=None)
          output_source = self.process
        else:
          self.process = subprocess.Popen(self.command, stdout=subprocess.PIPE)
          output_source = self.process.stdout
        last_run_time = datetime.datetime.now()
        for line in output_source:
          # If the process we're logging likes to repeat its output, we need to
          # look for the last line we saw before we start doing anything with
          # these lines anymore.
          if self.repeats_output_when_opened:
            if not should_log:
              if last_line == line:
                should_log = True
              continue
            elif skip_exception_line:
              # ignore the last line which caused UnicodeEncodeError
              skip_exception_line = False
              continue

          if self.output_fp:
            self.output_fp.write(line.decode("utf-8", "backslashreplace").rstrip())
            self.output_fp.write("\n")

          # Loop through all events we're watching for, to see if they occur on
          # this line. If they do, update the fact that we've seen this event.
          for event in self.events:
            if self.looking:
              event.ScanForEvent(line, lock=self.lock)
          last_line = line
      except UnicodeEncodeError:
        logging.exception("UnicodeEncodeError on running logger process")
        skip_exception_line = True
      except:  # pylint:disable=bare-except
        logging.exception("Exception encountered running process")
      finally:
        if pexpect:
          self.process.terminate()
        else:
          self.process.send_signal(signal.SIGTERM)
        should_log = False
      self.lock.acquire()

    self.lock.release()
    if pexpect:
      if self.process.exitstatus is not None:
        logging.info("Process finished - exit code %d", self.process.exitstatus)
      else:
        logging.info("Process finished - signal code %d",
                     self.process.signalstatus)
    else:
      if self.process.returncode is not None:
        logging.info("Process finished - return code %d",
                     self.process.returncode)
      else:
        logging.info("Process finished - no return code")

  def StopLogging(self):
    if self.process:
      self.lock.acquire()
      self.restart_process = False
      self.lock.release()

      if pexpect:
        self.process.kill(signal.SIGHUP)
        self.process.kill(signal.SIGINT)
      else:
        self.process.send_signal(signal.SIGTERM)


class Device(object):

  SECONDS_TO_SLEEP_DURING_ROOT = 0.5

  def __init__(self, serial_number, output_root, test_events, expected_result):
    """Responsible for monitoring a specific device, and pulling files from it.

    The actual work of the constructor will be handled asynchronously, you must
    call WaitForTasks() before using the device.

    Args:
      serial_number: The device serial number.
      output_root: The directory where to output log files/anything pulled from
          the device.
      test_events: The events (with conditions) that come from the StressTest
          that should be evaluated at every iteration, along with a list of
          actions to take when one of these events occur. For example, if there
          have not been any detected hotword triggers, a bugreport can be
          generated.
      expected_result: Expected event count to pass the test.
    """
    self.serial_number = serial_number
    self.output_root = output_root
    self.cmd_string_replacements = {}
    self.iteration = 0
    self.cmd_string_replacements["iteration"] = 0
    self.cmd_string_replacements["serial_number"] = serial_number
    self.cmd_string_replacements["output_root"] = output_root
    self.name = None
    self.process_loggers = []
    self.event_log = stress_test_pb2.EventLog()
    self.cnt_per_iteration = expected_result

    # Prepare the work queue, and offload the rest of the init into it.
    self.work_queue = queue.Queue()
    self.worker = threading.Thread(target=QueueWorker, args=[self.work_queue])
    self.worker.daemon = True
    self.worker.name = self.name
    self.worker.start()
    self.abort_requested = False
    self.remove_device = False
    self.test_events = test_events

    self.work_queue.put(self.__init_async__)

  def __init_async__(self):
    # Get the device type, and append it to the serial number.
    self.device_type = self.Command(["shell", "getprop",
                                     "ro.product.name"]).strip().decode("utf-8")
    self.name = "%s_%s" % (self.device_type, self.serial_number)
    self.worker.name = self.name
    self.cmd_string_replacements["device"] = self.name
    logging.info("Setting up device %s", self.name)

    config = stress_test_common.LoadDeviceConfig(self.device_type,
                                                 self.serial_number)

    # Get the device ready.
    self.Root()

    # Run any setup commands.
    for cmd in config.setup_command:
      result = self.Command(
          shlex.split(cmd % self.cmd_string_replacements)).strip()
      if result:
        for line in result.splitlines():
          logging.info(line)

    self.files_to_move = config.file_to_move

    self.event_names = set([event.name for event in config.event])
    self.event_counter = {name: 0 for name in self.event_names}
    self.iterations_since_event = {name: 0 for name in self.event_names}

    for file_to_watch in config.file_to_watch:
      # Are there any events that match up with this file?
      events = [x for x in config.event if x.source == file_to_watch.source]

      if file_to_watch.source == "LOGCAT":
        command = [
            "adb", "-s", self.serial_number, "logcat", "-v", "usec", ""
        ]
        command.extend(["%s:S" % tag for tag in config.tag_to_suppress])
        name = "logcat_" + self.serial_number
      else:
        command = [
            "adb", "-s", self.serial_number, "shell",
            "while : ; do cat %s 2>&1; done" % file_to_watch.source
        ]
        name = "%s_%s" % (os.path.basename(
            file_to_watch.source), self.serial_number)

      process_logger = ProcessLogger(
          name, command, os.path.join(
              self.output_root,
              file_to_watch.destination % self.cmd_string_replacements),
          events, True, file_to_watch.repeats_output_on_open)
      self.process_loggers.append(process_logger)
      process_logger.start()

    # Add any of the background processes.
    for daemon_process in config.daemon_process:
      # Are there any events that match up with this file?
      events = [x for x in config.event if x.source == daemon_process.name]
      command = shlex.split(
          daemon_process.command % self.cmd_string_replacements)
      if daemon_process.destination:
        output = os.path.join(
            self.output_root,
            daemon_process.destination % self.cmd_string_replacements)
      else:
        output = None
      name = "%s_%s" % (daemon_process.name, self.serial_number)
      process_logger = ProcessLogger(name, command, output, events,
                                     daemon_process.restart,
                                     daemon_process.repeats_output_on_open)
      self.process_loggers.append(process_logger)
      process_logger.start()

    # Build up the list of events we can actually process.
    self.__UpdateEventCounters(number_of_iterations=0)
    test_events = self.test_events
    self.test_events = []
    for event in test_events:
      try:
        eval(event.condition,  # pylint:disable=eval-used
             {"__builtins__": None}, self.__ValuesInEval())
        self.test_events.append(event)
      except Exception as err:  # pylint:disable=broad-except
        logging.error("Test event %s is not compatible with %s", event.name,
                      self.name)
        logging.error(str(err))
    # Make sure that device specific events don't have conditions.
    self.device_events = []
    for event in config.test_event:
      if not event.name:
        logging.error("Device %s test event is missing a name", self.name)
        continue
      if event.condition:
        self.test_events.append(event)
      else:
        self.device_events.append(event)

  def StartLookingForEvents(self):
    """Starts all child ProcessLoggers to start looking for events."""
    for process_logger in self.process_loggers:
      process_logger.looking = True

  def __ValuesInEval(self):
    values_in_eval = {key: value for key, value
                      in list(self.event_counter.items())}
    for key, value in list(self.iterations_since_event.items()):
      values_in_eval["iterations_since_%s" % key] = value
    return values_in_eval

  def __GetExpectedEventCount(self, event):
    if event == "logcat_iteration":
      return -1
    try:
      event_cnt = getattr(self.cnt_per_iteration, event)
    except AttributeError:
      event_cnt = -1
      logging.exception("%s is not an attribute of expected_result", event)
    return event_cnt

  def __UpdateEventCounters(self, number_of_iterations=1):
    # Update the event counters
    visited_events = set()
    error_log = []
    for process_logger in self.process_loggers:
      events = process_logger.GetEventCountsSinceLastCall()
      for event, count in list(events.items()):
        # Print log when there is any missed event
        expected_count = self.__GetExpectedEventCount(event)

        if expected_count > 0:
          if count > expected_count * number_of_iterations:
            logging.info(
                "[STRESS_TEST] In iteration %d, got duplicated %s : %d",
                self.iteration, self.name, count)
            logging.info("[STRESS_TEST] Will count only : %d",
                         expected_count * number_of_iterations)
            count = expected_count * number_of_iterations

        if count:
          self.event_counter[event] += count
          visited_events.add(event)

        if expected_count >= 0:
          if expected_count * number_of_iterations != count:
            error_log.append(
                _SUMMARY_COL_FORMATT %
                (event, count, expected_count * number_of_iterations))

    # Go clear all the events that weren't consecutive.
    for event in self.iterations_since_event:
      if event in visited_events:
        self.iterations_since_event[event] = 0
      else:
        self.iterations_since_event[event] += number_of_iterations

    if error_log:
      logging.info(_SUMMARY_LINES)
      logging.info(" iteration %d : Something wrong in %s.",
                   self.iteration, self.name)
      logging.info(_SUMMARY_LINES)
      logging.info(_SUMMARY_COLUMNS)
      logging.info(_SUMMARY_LINES)
      for line in error_log:
        logging.info(line)
      logging.info(_SUMMARY_LINES)

  def ProcessEvents(self):
    """Updates the event_counter and iterations_since_event maps."""
    self.work_queue.put(self.__ProcessEventsAsync)

  def __ProcessEventsAsync(self):
    # Move any files to the local machine that should be moved.
    if self.files_to_move:
      for file_to_move in self.files_to_move:
        try:
          self.Command(["pull", file_to_move.source, file_to_move.destination])
        except:  # pylint:disable=bare-except
          logging.exception("Failed to pull %s", file_to_move.source)

    self.__UpdateEventCounters()

    for event in self.test_events:
      if eval(event.condition,  # pylint:disable=eval-used
              {"__builtins__": None}, self.__ValuesInEval()):
        logging.info("Condition has been met for event '%s'", event.name)
        # Write the updated event log.
        event_log_details = self.event_log.event.add()
        event_log_details.iteration = self.iteration
        event_log_details.name = event.name
        with open(os.path.join(self.output_root,
                               "%s_event_log.ascii_proto" % self.name),
                  "w") as fp:
          text_format.PrintMessage(self.event_log, fp)

        # Do whatever other actions that are part of the event.
        self.__ProcessEventActionQueue(event)

        # Run any device specific actions for this event.
        for device_event in self.device_events:
          if device_event.name == event.name:
            self.__ProcessEventActionQueue(device_event)

    # Set up the next iteration.
    self.iteration += 1
    self.cmd_string_replacements["iteration"] = self.iteration

  def __ProcessEventActionQueue(self, event):
    bugreport = None
    for action in event.action:
      if action == "BUGREPORT":
        bugreport = self.TakeBugReport()
      elif action.startswith("DUMPSYS "):
        self.CaptureDumpsys(action[action.find(" ") + 1:])
      elif action == "NOTIFY":
        SendNotificationEmail(
            "%s had event '%s' occur" % (self.name, event.name),
            "\n".join(["Current Summary:"] + self.GetSummaryLines()), bugreport)
      elif action == "REMOVE_DEVICE":
        logging.info("Removing %s from the test", self.serial_number)
        self.remove_device = True
      elif action == "ABORT":
        logging.info("Abort requested")
        self.abort_requested = True
      else:
        action %= self.cmd_string_replacements
        logging.info("Running command %s on %s", action, self.name)
        result = self.Command(shlex.split(action)).strip()
        if result:
          for line in result.splitlines():
            logging.info(line)

  def Root(self):
    self.Command(["root"])
    time.sleep(Device.SECONDS_TO_SLEEP_DURING_ROOT)
    self.Command(["wait-for-device"])
    time.sleep(Device.SECONDS_TO_SLEEP_DURING_ROOT)

  def Stop(self):
    """Stops all file loggers attached to this device."""
    for process_logger in self.process_loggers:
      process_logger.StopLogging()
    self.process_loggers = []

  def Join(self):
    for process_logger in self.process_loggers:
      process_logger.join()
    self.WaitForTasks()

  def AsyncCommand(self, command, log_output=False):
    self.work_queue.put(
        lambda: self.__AsyncCommand(command, log_output=log_output))

  def __AsyncCommand(self, command, log_output=False):
    result = self.Command(command).strip()
    if result and log_output:
      for line in result.splitlines():
        logging.info(line.decode("utf-8"))

  def Command(self, command):
    """Runs the provided command on this device."""
    if command[0] in {"bugreport", "root", "wait-for-device", "shell",
                      "logcat"}:
      return subprocess.check_output(
          ["adb", "-s", self.serial_number] + command)
    elif command[0] == "DUMPSYS":
      self.CaptureDumpsys(command[1])
      return ""
    elif command[0] == "pull":
      try:
        files = subprocess.check_output(
            ["adb", "-s", self.serial_number, "shell", "ls", command[1]]
        ).strip().splitlines()
      except subprocess.CalledProcessError:
        return ""
      if len(files) == 1 and "No such file or directory" in files[0]:
        return ""
      for source_file in files:
        destination = os.path.join(self.output_root,
                                   command[2] % self.cmd_string_replacements)
        stress_test_common.MakeDirsIfNeeded(os.path.dirname(destination))
        logging.info("Moving %s from %s to %s", source_file, self.name,
                     destination)
        subprocess.check_output(["adb", "-s", self.serial_number, "pull",
                                 source_file, destination])
        if FLAGS.delete_data_dir:
          subprocess.check_output([
              "adb", "-s", self.serial_number, "shell", "rm", "-rf", source_file
          ])
        return ""
    else:
      return subprocess.check_output(command)

  def TakeBugReport(self):
    logging.info("Capturing bugreport on %s", self.name)
    bugreport = os.path.join(self.output_root,
                             "%s_bugreport_iteration_%06d.zip" %
                             (self.name, self.iteration))
    sdk = int(self.Command(
        ["shell", "getprop", "ro.build.version.sdk"]).strip())
    if sdk >= 24:  # SDK 24 = Android N
      with open(bugreport, "w") as bugreport_fp:
        bugreport_fp.write(self.Command(["bugreport", bugreport]))
    else:
      bugreport_txt = os.path.join(self.output_root,
                                   "%s_bugreport_iteration_%06d.txt" %
                                   (self.name, self.iteration))
      with open(bugreport_txt, "w") as bugreport_fp:
        bugreport_fp.write(self.Command(["bugreport"]))
      self.Command(["zip", bugreport, bugreport_txt])

    self.Command(["pull", "/data/anr/traces.txt",
                  "%s_traces_iteration_%06d.txt" % (self.name, self.iteration)])
    self.Command(["pull", "/data/anr/traces.txt.bugreport",
                  "%s_traces_iteration_%06d.txt.bugreport" % (self.name,
                                                              self.iteration)])
    return bugreport

  def CaptureDumpsys(self, dumpsys_unit):
    logging.info("Taking dumpsys %s on %s", dumpsys_unit, self.name)
    stress_test_common.MakeDirsIfNeeded(os.path.join(self.output_root,
                                                     self.name))
    with open(os.path.join(self.output_root, self.name,
                           "%s_%06d.txt" % (dumpsys_unit, self.iteration)),
              "w") as dumpsys_fp:
      dumpsys_fp.write(self.Command(["shell", "dumpsys", dumpsys_unit]))

  def WaitForTasks(self):
    self.work_queue.join()

  def GetSummaryLines(self):
    lines = [
        "Device {}".format(self.name),
        _SUMMARY_LINES, _SUMMARY_COLUMNS, _SUMMARY_LINES
    ]
    for event, count in sorted(self.event_counter.items()):
      lines.append(_SUMMARY_COL_FORMATT % (
          event, count, self.iterations_since_event[event]))
    lines.append(_SUMMARY_LINES)
    return lines


def RunAsyncCommand(devices, command):
  """Helper function for running async commands on many devices."""
  for device in devices:
    device.AsyncCommand(command)
  for device in devices:
    device.WaitForTasks()


class StressTest(object):
  """Manages dispatching commands to devices/playing audio and events."""

  def __init__(self, output_root, test_name):
    self.output_root = output_root
    self.devices = []
    self.test_name = test_name
    config = stress_test_pb2.StressTestConfig()
    config_contents = stress_test_common.GetResourceContents(
        os.path.join(stress_test_common.RESOURCE_DIR,
                     "stress_test.%s.ascii_proto" % test_name))
    text_format.Merge(config_contents, config)
    self.events = config.event
    self.setup_commands = config.setup_command
    self.steps = config.step
    self.audio_tempfiles = {}
    self.uuid = str(uuid.uuid4())
    self.expected_result = None
    self.iteration = 0
    if config.expected_result:
      self.expected_result = config.expected_result[0]

    # Place all the audio files into temp files.
    for step in self.steps:
      if step.audio_file and step.audio_file not in self.audio_tempfiles:
        # We can't delete the temp file on windows, since it gets nuked too
        # early.
        audio_tempfile = tempfile.NamedTemporaryFile(
            delete=(platform.system() != "Windows"),
            dir="." if platform.system().startswith("CYGWIN") else None
        )
        if platform.system().startswith("CYGWIN"):
          audio_tempfile.name = os.path.basename(audio_tempfile.name)
        self.audio_tempfiles[step.audio_file] = audio_tempfile
        if FLAGS.use_sox:
          # Write out the raw PCM samples as a wave file.
          audio_tempfile.write(
              stress_test_common.GetResourceContents(step.audio_file))
        else:
          # Make a temporary wave file for playout if we can't use sox.
          wavefile = wave.open(audio_tempfile, "wb")
          if step.audio_file_sample_rate <= 0:
            step.audio_file_sample_rate = 16000
          wavefile.setframerate(step.audio_file_sample_rate)
          if step.audio_file_num_channels <= 0:
            step.audio_file_num_channels = 1
          wavefile.setnchannels(step.audio_file_num_channels)
          if not step.audio_file_format:
            wavefile.setsampwidth(2)
          elif step.audio_file_format == "s8":
            wavefile.setsampwidth(1)
          elif step.audio_file_format == "s16":
            wavefile.setsampwidth(2)
          elif step.audio_file_format == "s32":
            wavefile.setsampwidth(4)
          else:
            raise RuntimeError(
                "Unsupported wave file format for %s" % step.audio_file)
          wavefile.writeframes(stress_test_common.GetResourceContents(
              step.audio_file))
          wavefile.close()
        audio_tempfile.flush()

        if platform.system() == "Windows":
          audio_tempfile.close()

    # Create all the devices that are attached to this machine.
    for serial_number in self.GetActiveSerialNumbers():
      self.devices.append(
          Device(serial_number, output_root, self.events, self.expected_result))
    if not self.devices:
      raise app.UsageError("No devices connected")

    self.devices.sort(key=lambda x: x.name)

    # Make sure every device is done with their work for setup.
    for device in self.devices:
      device.WaitForTasks()

    # Write out the info meta-data proto. Useful for doing analysis of the logs
    # after the stress test has completed.
    stress_test_info = stress_test_pb2.StressTestInfo()
    stress_test_info.test_name = self.test_name
    stress_test_info.test_description = config.description
    stress_test_info.uuid = self.uuid
    for device in self.devices:
      device_pb = stress_test_info.device.add()
      device_pb.device_type = device.device_type
      device_pb.serial_number = device.serial_number

    text_format.PrintMessage(stress_test_info, open(os.path.join(
        self.output_root, "stress_test_info.ascii_proto"), "w"))

  def GetActiveSerialNumbers(self):
    serial_numbers = []
    for line in sorted(
        subprocess.check_output(["adb", "devices"]).splitlines()):
      if line.endswith(b"device"):
        serial_number = line.split()[0].strip()
        if FLAGS.devices and serial_number not in FLAGS.devices:
          continue
        serial_numbers.append(serial_number.decode("utf-8"))
    return serial_numbers

  def Start(self):
    logging.info("Waiting for devices to settle")
    time.sleep(5)
    # Make a copy of the device list, as we'll be modifying this actual list.
    devices = list(self.devices)
    dropped_devices = []

    # If we have any setup commands, run them.
    for command in self.setup_commands:
      logging.info("Running command %s", command)
      # Can't use the async command helper function since we need to get at
      # the device cmd_string_replacements.
      for device in devices:
        device.AsyncCommand(
            shlex.split(command % device.cmd_string_replacements),
            log_output=True)
      for device in devices:
        device.WaitForTasks()

    for device in devices:
      device.StartLookingForEvents()
      device.AsyncCommand(["shell", "log", "-t", "STRESS_TEST",
                           "Starting {%s} TZ=$(getprop persist.sys.timezone) "
                           "YEAR=$(date +%%Y)" % self.uuid], True)
    self.iteration = 0
    while True:
      logging.info("Starting iteration %d", self.iteration)
      # Perform all the actions specified in the test.
      RunAsyncCommand(devices, [
          "shell", "log", "-t", "STRESS_TEST",
          "Performing iteration %d $(head -n 3 "
          "/proc/timer_list | tail -n 1)" % self.iteration
      ])

      for step in self.steps:
        if step.delay_before:
          logging.info("Waiting for %.2f seconds", step.delay_before)
          time.sleep(step.delay_before)

        if step.audio_file:
          logging.info("Playing %s", step.audio_file)
          RunAsyncCommand(devices, ["shell", "log", "-t", "STRESS_TEST",
                                    "Playing %s" % step.audio_file])

          if FLAGS.use_sox:
            subprocess.check_call(["sox", "-q",
                                   self.audio_tempfiles[step.audio_file].name,
                                   "-d"])
          elif platform.system() == "Windows":
            import winsound  # pylint:disable=g-import-not-at-top
            winsound.PlaySound(self.audio_tempfiles[step.audio_file].name,
                               winsound.SND_FILENAME | winsound.SND_NODEFAULT)
          else:
            raise app.RuntimeError("Unsupported platform for audio playback")

        if step.command:
          logging.info("Running command %s", step.command)
          # Can't use the async command helper function since we need to get at
          # the device cmd_string_replacements.
          for device in devices:
            device.AsyncCommand(
                shlex.split(step.command % device.cmd_string_replacements),
                log_output=True)
          for device in devices:
            device.WaitForTasks()

        if step.delay_after:
          logging.info("Waiting for %.2f seconds", step.delay_after)
          time.sleep(step.delay_after)

      RunAsyncCommand(devices, [
          "shell", "log", "-t", "STRESS_TEST",
          "Iteration %d complete $(head -n 3 "
          "/proc/timer_list | tail -n 1)" % self.iteration
      ])
      self.iteration += 1

      # TODO(somebody): Sometimes the logcat seems to get stuck and buffers for
      # a bit. This throws off the event counts, so we should probably add some
      # synchronization rules before we trigger any events.

      # Go through each device, update the event counter, and see if we need to
      # trigger any events.
      devices_to_remove = []
      abort_requested = False
      active_devices = self.GetActiveSerialNumbers()
      for device in devices:
        if device.serial_number in active_devices:
          device.ProcessEvents()
        else:
          logging.error("Dropped device %s", device.name)
          SendNotificationEmail(
              "Dropped device %s" % device.name,
              "Device %s is not longer present in the system" % device.name)
          dropped_devices.append(device)
          devices_to_remove.append(device)

      # Check to see if any of the dropped devices have come back. If yes, grab
      # a bug report.
      for device in dropped_devices:
        if device.serial_number in active_devices:
          logging.info("Device %s reappeared", device.name)
          device.Root()
          device.TakeBugReport()

      dropped_devices = [d for d in dropped_devices
                         if d.serial_number not in active_devices]

      for device in devices:
        device.WaitForTasks()
        if device.remove_device:
          devices_to_remove.append(device)
        if device.abort_requested:
          abort_requested = True

      # Remove devices from our list of things to monitor if they've been marked
      # for deletion.
      if devices_to_remove:
        for device in devices_to_remove:
          device.Stop()
        devices = [d for d in devices if d not in devices_to_remove]

      # Print out the iteration summary.
      if self.iteration % FLAGS.print_summary_every_n == 0:
        for line in self.GetSummaryLines():
          logging.info(line)

      # See if we need to break out of the outer loop.
      if abort_requested or not devices:
        break
      if FLAGS.num_iterations:
        if self.iteration >= FLAGS.num_iterations:
          logging.info("Completed full iteration : %d", self.iteration)
          break
    SendNotificationEmail(
        "Stress test %s completed" % (FLAGS.test_name),
        "\n".join(["Summary:"] + self.GetSummaryLines()))

  def Stop(self):
    logging.debug("Stopping devices")
    for device in self.devices:
      device.Stop()
    for device in self.devices:
      device.Join()

  def GetSummaryLines(self):
    lines = [
        _SUMMARY_LINES,
        "Conducted %d iterations out of %d" %
        (self.iteration, FLAGS.num_iterations),
        _SUMMARY_LINES
    ]
    for device in self.devices:
      lines.extend(device.GetSummaryLines())
    lines.append(_SUMMARY_LINES)
    return lines


def main(unused_argv):
  # Check to make sure that there are no other instances of ADB running - if
  # there are, print a warning and wait a bit for them to see it and decide if
  # they want to keep running, knowing that logs may be invalid.
  try:
    if "adb" in subprocess.check_output(["ps", "-ale"]).decode("utf-8"):
      print("It looks like there are other instances of adb running. If these "
            "other instances are also cating log files, you will not be "
            "capturing everything in this stress test (so logs will be "
            "invalid).")
      print("Continuing in 3...", end=" ")
      sys.stdout.flush()
      for i in [2, 1, 0]:
        time.sleep(1)
        if i:
          print("%d..." % i, end=" ")
        else:
          print("")
        sys.stdout.flush()
  except OSError:
    print("Unexpected error:", sys.exc_info()[0])
    if sys.platform.startswith("win"):
      pass
    else:
      raise

  # Make the base output directory.
  output_root = os.path.join(FLAGS.output_root, "%s_%s" % (
      FLAGS.test_name, datetime.datetime.now().strftime("%Y%m%d_%H%M%S")))
  # output_root = os.path.join(FLAGS.output_root, FLAGS.test_name)
  stress_test_common.MakeDirsIfNeeded(output_root)

  # Set up logging.
  formatter = logging.Formatter(
      "%(levelname)-1.1s %(asctime)s [%(threadName)-16.16s] %(message)s")
  root_logger = logging.getLogger()
  root_logger.setLevel(logging.INFO)
  root_logger.setLevel(logging.DEBUG)

  file_handler = logging.FileHandler(os.path.join(output_root,
                                                  "stress_test.log"))
  file_handler.setFormatter(formatter)
  root_logger.addHandler(file_handler)

  console_handler = logging.StreamHandler()
  console_handler.setFormatter(formatter)
  root_logger.addHandler(console_handler)

  stress_test = StressTest(output_root, FLAGS.test_name)
  try:
    stress_test.Start()
  finally:
    logging.info("Stopping device logging threads")
    stress_test.Stop()
    for line in stress_test.GetSummaryLines():
      logging.info(line)
    if FLAGS.delete_data_dir:
      print("Deleting Data Dir")
      subprocess.check_output(["rm", "-r", "-f", output_root])


if __name__ == "__main__":
  app.run(main)