• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2016 Google Inc.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import base64
16import concurrent.futures
17import datetime
18import errno
19import inspect
20import io
21import logging
22import os
23import platform
24import random
25import re
26import shlex
27import signal
28import string
29import subprocess
30import threading
31import time
32import traceback
33from typing import Literal, Tuple, overload
34
35import portpicker
36
37# File name length is limited to 255 chars on some OS, so we need to make sure
38# the file names we output fits within the limit.
39MAX_FILENAME_LEN = 255
40# Number of times to retry to get available port
41MAX_PORT_ALLOCATION_RETRY = 50
42
43ascii_letters_and_digits = string.ascii_letters + string.digits
44valid_filename_chars = f'-_.{ascii_letters_and_digits}'
45
46GMT_to_olson = {
47    'GMT-9': 'America/Anchorage',
48    'GMT-8': 'US/Pacific',
49    'GMT-7': 'US/Mountain',
50    'GMT-6': 'US/Central',
51    'GMT-5': 'US/Eastern',
52    'GMT-4': 'America/Barbados',
53    'GMT-3': 'America/Buenos_Aires',
54    'GMT-2': 'Atlantic/South_Georgia',
55    'GMT-1': 'Atlantic/Azores',
56    'GMT+0': 'Africa/Casablanca',
57    'GMT+1': 'Europe/Amsterdam',
58    'GMT+2': 'Europe/Athens',
59    'GMT+3': 'Europe/Moscow',
60    'GMT+4': 'Asia/Baku',
61    'GMT+5': 'Asia/Oral',
62    'GMT+6': 'Asia/Almaty',
63    'GMT+7': 'Asia/Bangkok',
64    'GMT+8': 'Asia/Hong_Kong',
65    'GMT+9': 'Asia/Tokyo',
66    'GMT+10': 'Pacific/Guam',
67    'GMT+11': 'Pacific/Noumea',
68    'GMT+12': 'Pacific/Fiji',
69    'GMT+13': 'Pacific/Tongatapu',
70    'GMT-11': 'Pacific/Midway',
71    'GMT-10': 'Pacific/Honolulu',
72}
73
74
75class Error(Exception):
76  """Raised when an error occurs in a util"""
77
78
79def abs_path(path):
80  """Resolve the '.' and '~' in a path to get the absolute path.
81
82  Args:
83    path: The path to expand.
84
85  Returns:
86    The absolute path of the input path.
87  """
88  return os.path.abspath(os.path.expanduser(path))
89
90
91def create_dir(path):
92  """Creates a directory if it does not exist already.
93
94  Args:
95    path: The path of the directory to create.
96  """
97  full_path = abs_path(path)
98  if not os.path.exists(full_path):
99    try:
100      os.makedirs(full_path)
101    except OSError as e:
102      # ignore the error for dir already exist.
103      if e.errno != errno.EEXIST:
104        raise
105
106
107def create_alias(target_path, alias_path):
108  """Creates an alias at 'alias_path' pointing to the file 'target_path'.
109
110  On Unix, this is implemented via symlink. On Windows, this is done by
111  creating a Windows shortcut file.
112
113  Args:
114    target_path: Destination path that the alias should point to.
115    alias_path: Path at which to create the new alias.
116  """
117  if platform.system() == 'Windows' and not alias_path.endswith('.lnk'):
118    alias_path += '.lnk'
119  if os.path.lexists(alias_path):
120    os.remove(alias_path)
121  if platform.system() == 'Windows':
122    from win32com import client
123
124    shell = client.Dispatch('WScript.Shell')
125    shortcut = shell.CreateShortCut(alias_path)
126    shortcut.Targetpath = target_path
127    shortcut.save()
128  else:
129    os.symlink(target_path, alias_path)
130
131
132def get_current_epoch_time():
133  """Current epoch time in milliseconds.
134
135  Returns:
136    An integer representing the current epoch time in milliseconds.
137  """
138  return int(round(time.time() * 1000))
139
140
141def get_current_human_time():
142  """Returns the current time in human readable format.
143
144  Returns:
145    The current time stamp in Month-Day-Year Hour:Min:Sec format.
146  """
147  return time.strftime('%m-%d-%Y %H:%M:%S ')
148
149
150def epoch_to_human_time(epoch_time):
151  """Converts an epoch timestamp to human readable time.
152
153  This essentially converts an output of get_current_epoch_time to an output
154  of get_current_human_time
155
156  Args:
157    epoch_time: An integer representing an epoch timestamp in milliseconds.
158
159  Returns:
160    A time string representing the input time.
161    None if input param is invalid.
162  """
163  if isinstance(epoch_time, int):
164    try:
165      d = datetime.datetime.fromtimestamp(epoch_time / 1000)
166      return d.strftime('%m-%d-%Y %H:%M:%S ')
167    except ValueError:
168      return None
169
170
171def get_timezone_olson_id():
172  """Return the Olson ID of the local (non-DST) timezone.
173
174  Returns:
175    A string representing one of the Olson IDs of the local (non-DST)
176    timezone.
177  """
178  tzoffset = int(time.timezone / 3600)
179  if tzoffset <= 0:
180    gmt = f'GMT+{-tzoffset}'
181  else:
182    gmt = f'GMT-{tzoffset}'
183  return GMT_to_olson[gmt]
184
185
186def find_files(paths, file_predicate):
187  """Locate files whose names and extensions match the given predicate in
188  the specified directories.
189
190  Args:
191    paths: A list of directory paths where to find the files.
192    file_predicate: A function that returns True if the file name and
193      extension are desired.
194
195  Returns:
196    A list of files that match the predicate.
197  """
198  file_list = []
199  for path in paths:
200    p = abs_path(path)
201    for dirPath, _, fileList in os.walk(p):
202      for fname in fileList:
203        name, ext = os.path.splitext(fname)
204        if file_predicate(name, ext):
205          file_list.append((dirPath, name, ext))
206  return file_list
207
208
209def load_file_to_base64_str(f_path):
210  """Loads the content of a file into a base64 string.
211
212  Args:
213    f_path: full path to the file including the file name.
214
215  Returns:
216    A base64 string representing the content of the file in utf-8 encoding.
217  """
218  path = abs_path(f_path)
219  with io.open(path, 'rb') as f:
220    f_bytes = f.read()
221    base64_str = base64.b64encode(f_bytes).decode('utf-8')
222    return base64_str
223
224
225def find_field(item_list, cond, comparator, target_field):
226  """Finds the value of a field in a dict object that satisfies certain
227  conditions.
228
229  Args:
230    item_list: A list of dict objects.
231    cond: A param that defines the condition.
232    comparator: A function that checks if an dict satisfies the condition.
233    target_field: Name of the field whose value to be returned if an item
234      satisfies the condition.
235
236  Returns:
237    Target value or None if no item satisfies the condition.
238  """
239  for item in item_list:
240    if comparator(item, cond) and target_field in item:
241      return item[target_field]
242  return None
243
244
245def rand_ascii_str(length):
246  """Generates a random string of specified length, composed of ascii letters
247  and digits.
248
249  Args:
250    length: The number of characters in the string.
251
252  Returns:
253    The random string generated.
254  """
255  letters = [random.choice(ascii_letters_and_digits) for _ in range(length)]
256  return ''.join(letters)
257
258
259# Thead/Process related functions.
260def _collect_process_tree(starting_pid):
261  """Collects PID list of the descendant processes from the given PID.
262
263  This function only available on Unix like system.
264
265  Args:
266    starting_pid: The PID to start recursively traverse.
267
268  Returns:
269    A list of pid of the descendant processes.
270  """
271  ret = []
272  stack = [starting_pid]
273
274  while stack:
275    pid = stack.pop()
276    if platform.system() == 'Darwin':
277      command = ['pgrep', '-P', str(pid)]
278    else:
279      command = [
280          'ps',
281          '-o',
282          'pid',
283          '--ppid',
284          str(pid),
285          '--noheaders',
286      ]
287    try:
288      ps_results = subprocess.check_output(command).decode().strip()
289    except subprocess.CalledProcessError:
290      # Ignore if there is not child process.
291      continue
292
293    children_pid_list = [int(p.strip()) for p in ps_results.split('\n')]
294    stack.extend(children_pid_list)
295    ret.extend(children_pid_list)
296
297  return ret
298
299
300def _kill_process_tree(proc):
301  """Kills the subprocess and its descendants."""
302  if os.name == 'nt':
303    # The taskkill command with "/T" option ends the specified process and any
304    # child processes started by it:
305    # https://docs.microsoft.com/en-us/windows-server/administration/windows-commands/taskkill
306    subprocess.check_output(
307        [
308            'taskkill',
309            '/F',
310            '/T',
311            '/PID',
312            str(proc.pid),
313        ]
314    )
315    return
316
317  failed = []
318  for child_pid in _collect_process_tree(proc.pid):
319    try:
320      os.kill(child_pid, signal.SIGTERM)
321    except Exception:  # pylint: disable=broad-except
322      failed.append(child_pid)
323      logging.exception('Failed to kill standing subprocess %d', child_pid)
324
325  try:
326    proc.kill()
327  except Exception:  # pylint: disable=broad-except
328    failed.append(proc.pid)
329    logging.exception('Failed to kill standing subprocess %d', proc.pid)
330
331  if failed:
332    raise Error('Failed to kill standing subprocesses: %s' % failed)
333
334
335def concurrent_exec(func, param_list, max_workers=30, raise_on_exception=False):
336  """Executes a function with different parameters pseudo-concurrently.
337
338  This is basically a map function. Each element (should be an iterable) in
339  the param_list is unpacked and passed into the function. Due to Python's
340  GIL, there's no true concurrency. This is suited for IO-bound tasks.
341
342  Args:
343    func: The function that performs a task.
344    param_list: A list of iterables, each being a set of params to be
345      passed into the function.
346    max_workers: int, the number of workers to use for parallelizing the
347      tasks. By default, this is 30 workers.
348    raise_on_exception: bool, raises all of the task failures if any of the
349      tasks failed if `True`. By default, this is `False`.
350
351  Returns:
352    A list of return values from each function execution. If an execution
353    caused an exception, the exception object will be the corresponding
354    result.
355
356  Raises:
357    RuntimeError: If executing any of the tasks failed and
358      `raise_on_exception` is True.
359  """
360  with concurrent.futures.ThreadPoolExecutor(
361      max_workers=max_workers
362  ) as executor:
363    # Start the load operations and mark each future with its params
364    future_to_params = {executor.submit(func, *p): p for p in param_list}
365    return_vals = []
366    exceptions = []
367    for future in concurrent.futures.as_completed(future_to_params):
368      params = future_to_params[future]
369      try:
370        return_vals.append(future.result())
371      except Exception as exc:  # pylint: disable=broad-except
372        logging.exception(
373            '%s generated an exception: %s', params, traceback.format_exc()
374        )
375        return_vals.append(exc)
376        exceptions.append(exc)
377    if raise_on_exception and exceptions:
378      error_messages = []
379      for exception in exceptions:
380        error_messages.append(
381            ''.join(
382                traceback.format_exception(
383                    exception.__class__, exception, exception.__traceback__
384                )
385            )
386        )
387      raise RuntimeError('\n\n'.join(error_messages))
388    return return_vals
389
390
391# Provide hint for pytype checker to avoid the Union[bytes, str] case.
392@overload
393def run_command(
394    cmd,
395    stdout=...,
396    stderr=...,
397    shell=...,
398    timeout=...,
399    cwd=...,
400    env=...,
401    universal_newlines: Literal[False] = ...,
402) -> Tuple[int, bytes, bytes]:
403  ...
404
405
406@overload
407def run_command(
408    cmd,
409    stdout=...,
410    stderr=...,
411    shell=...,
412    timeout=...,
413    cwd=...,
414    env=...,
415    universal_newlines: Literal[True] = ...,
416) -> Tuple[int, str, str]:
417  ...
418
419
420def run_command(
421    cmd,
422    stdout=None,
423    stderr=None,
424    shell=False,
425    timeout=None,
426    cwd=None,
427    env=None,
428    universal_newlines=False,
429):
430  """Runs a command in a subprocess.
431
432  This function is very similar to subprocess.check_output. The main
433  difference is that it returns the return code and std error output as well
434  as supporting a timeout parameter.
435
436  Args:
437    cmd: string or list of strings, the command to run.
438      See subprocess.Popen() documentation.
439    stdout: file handle, the file handle to write std out to. If None is
440      given, then subprocess.PIPE is used. See subprocess.Popen()
441      documentation.
442    stderr: file handle, the file handle to write std err to. If None is
443      given, then subprocess.PIPE is used. See subprocess.Popen()
444      documentation.
445    shell: bool, True to run this command through the system shell,
446      False to invoke it directly. See subprocess.Popen() docs.
447    timeout: float, the number of seconds to wait before timing out.
448      If not specified, no timeout takes effect.
449    cwd: string, the path to change the child's current directory to before
450      it is executed. Note that this directory is not considered when
451      searching the executable, so you can't specify the program's path
452      relative to cwd.
453    env: dict, a mapping that defines the environment variables for the
454      new process. Default behavior is inheriting the current process'
455      environment.
456    universal_newlines: bool, True to open file objects in text mode, False in
457      binary mode.
458
459  Returns:
460    A 3-tuple of the consisting of the return code, the std output, and the
461      std error.
462
463  Raises:
464    subprocess.TimeoutExpired: The command timed out.
465  """
466  if stdout is None:
467    stdout = subprocess.PIPE
468  if stderr is None:
469    stderr = subprocess.PIPE
470  process = subprocess.Popen(
471      cmd,
472      stdout=stdout,
473      stderr=stderr,
474      shell=shell,
475      cwd=cwd,
476      env=env,
477      universal_newlines=universal_newlines,
478  )
479  timer = None
480  timer_triggered = threading.Event()
481  if timeout and timeout > 0:
482    # The wait method on process will hang when used with PIPEs with large
483    # outputs, so use a timer thread instead.
484
485    def timeout_expired():
486      timer_triggered.set()
487      process.terminate()
488
489    timer = threading.Timer(timeout, timeout_expired)
490    timer.start()
491  # If the command takes longer than the timeout, then the timer thread
492  # will kill the subprocess, which will make it terminate.
493  out, err = process.communicate()
494  if timer is not None:
495    timer.cancel()
496  if timer_triggered.is_set():
497    raise subprocess.TimeoutExpired(
498        cmd=cmd, timeout=timeout, output=out, stderr=err
499    )
500  logging.debug(
501      'cmd: %s, stdout: %s, stderr: %s, ret: %s',
502      cli_cmd_to_string(cmd),
503      out,
504      err,
505      process.returncode,
506  )
507  return process.returncode, out, err
508
509
510def start_standing_subprocess(
511    cmd,
512    shell=False,
513    env=None,
514    stdout=subprocess.PIPE,
515    stderr=subprocess.PIPE,
516):
517  """Starts a long-running subprocess.
518
519  This is not a blocking call and the subprocess started by it should be
520  explicitly terminated with stop_standing_subprocess.
521
522  For short-running commands, you should use subprocess.check_call, which
523  blocks.
524
525  Args:
526    cmd: string, the command to start the subprocess with.
527    shell: bool, True to run this command through the system shell,
528      False to invoke it directly. See subprocess.Popen() docs.
529    env: dict, a custom environment to run the standing subprocess. If not
530      specified, inherits the current environment. See subprocess.Popen()
531      docs.
532    stdout: None, subprocess.PIPE, subprocess.DEVNULL, an existing file
533      descriptor, or an existing file object. See subprocess.Popen() docs.
534    stderr: None, subprocess.PIPE, subprocess.DEVNULL, an existing file
535      descriptor, or an existing file object. See subprocess.Popen() docs.
536
537  Returns:
538    The subprocess that was started.
539  """
540  logging.debug('Starting standing subprocess with: %s', cmd)
541  proc = subprocess.Popen(
542      cmd,
543      stdin=subprocess.PIPE,
544      stdout=stdout,
545      stderr=stderr,
546      shell=shell,
547      env=env,
548  )
549  # Leaving stdin open causes problems for input, e.g. breaking the
550  # code.inspect() shell (http://stackoverflow.com/a/25512460/1612937), so
551  # explicitly close it assuming it is not needed for standing subprocesses.
552  proc.stdin.close()
553  proc.stdin = None
554  logging.debug('Started standing subprocess %d', proc.pid)
555  return proc
556
557
558def stop_standing_subprocess(proc):
559  """Stops a subprocess started by start_standing_subprocess.
560
561  Before killing the process, we check if the process is running, if it has
562  terminated, Error is raised.
563
564  Catches and ignores the PermissionError which only happens on Macs.
565
566  Args:
567    proc: Subprocess to terminate.
568
569  Raises:
570    Error: if the subprocess could not be stopped.
571  """
572  logging.debug('Stopping standing subprocess %d', proc.pid)
573
574  _kill_process_tree(proc)
575
576  # Call wait and close pipes on the original Python object so we don't get
577  # runtime warnings.
578  if proc.stdout:
579    proc.stdout.close()
580  if proc.stderr:
581    proc.stderr.close()
582  proc.wait()
583  logging.debug('Stopped standing subprocess %d', proc.pid)
584
585
586def wait_for_standing_subprocess(proc, timeout=None):
587  """Waits for a subprocess started by start_standing_subprocess to finish
588  or times out.
589
590  Propagates the exception raised by the subprocess.wait(.) function.
591  The subprocess.TimeoutExpired exception is raised if the process timed-out
592  rather than terminating.
593
594  If no exception is raised: the subprocess terminated on its own. No need
595  to call stop_standing_subprocess() to kill it.
596
597  If an exception is raised: the subprocess is still alive - it did not
598  terminate. Either call stop_standing_subprocess() to kill it, or call
599  wait_for_standing_subprocess() to keep waiting for it to terminate on its
600  own.
601
602  If the corresponding subprocess command generates a large amount of output
603  and this method is called with a timeout value, then the command can hang
604  indefinitely. See http://go/pylib/subprocess.html#subprocess.Popen.wait
605
606  This function does not support Python 2.
607
608  Args:
609    p: Subprocess to wait for.
610    timeout: An integer number of seconds to wait before timing out.
611  """
612  proc.wait(timeout)
613
614
615def get_available_host_port():
616  """Gets a host port number available for adb forward.
617
618  DEPRECATED: This method is unreliable. Pass `tcp:0` to adb forward instead.
619
620  Returns:
621    An integer representing a port number on the host available for adb
622    forward.
623
624  Raises:
625    Error: when no port is found after MAX_PORT_ALLOCATION_RETRY times.
626  """
627  logging.warning(
628      'The method mobly.utils.get_available_host_port is deprecated because it '
629      'is unreliable. Pass "tcp:0" to adb forward instead.'
630  )
631
632  # Only import adb module if needed.
633  from mobly.controllers.android_device_lib import adb
634
635  port = portpicker.pick_unused_port()
636  if not adb.is_adb_available():
637    return port
638  for _ in range(MAX_PORT_ALLOCATION_RETRY):
639    # Make sure adb is not using this port so we don't accidentally
640    # interrupt ongoing runs by trying to bind to the port.
641    if port not in adb.list_occupied_adb_ports():
642      return port
643    port = portpicker.pick_unused_port()
644  raise Error(
645      'Failed to find available port after {} retries'.format(
646          MAX_PORT_ALLOCATION_RETRY
647      )
648  )
649
650
651def grep(regex, output):
652  """Similar to linux's `grep`, this returns the line in an output stream
653  that matches a given regex pattern.
654
655  It does not rely on the `grep` binary and is not sensitive to line endings,
656  so it can be used cross-platform.
657
658  Args:
659    regex: string, a regex that matches the expected pattern.
660    output: byte string, the raw output of the adb cmd.
661
662  Returns:
663    A list of strings, all of which are output lines that matches the
664    regex pattern.
665  """
666  lines = output.decode('utf-8').strip().splitlines()
667  results = []
668  for line in lines:
669    if re.search(regex, line):
670      results.append(line.strip())
671  return results
672
673
674def cli_cmd_to_string(args):
675  """Converts a cmd arg list to string.
676
677  Args:
678    args: list of strings, the arguments of a command.
679
680  Returns:
681    String representation of the command.
682  """
683  if isinstance(args, str):
684    # Return directly if it's already a string.
685    return args
686  return ' '.join([shlex.quote(arg) for arg in args])
687
688
689def get_settable_properties(cls):
690  """Gets the settable properties of a class.
691
692  Only returns the explicitly defined properties with setters.
693
694  Args:
695    cls: A class in Python.
696  """
697  results = []
698  for attr, value in vars(cls).items():
699    if isinstance(value, property) and value.fset is not None:
700      results.append(attr)
701  return results
702
703
704def find_subclasses_in_module(base_classes, module):
705  """Finds the subclasses of the given classes in the given module.
706
707  Args:
708    base_classes: list of classes, the base classes to look for the
709      subclasses of in the module.
710    module: module, the module to look for the subclasses in.
711
712  Returns:
713    A list of all of the subclasses found in the module.
714  """
715  subclasses = []
716  for _, module_member in module.__dict__.items():
717    if inspect.isclass(module_member):
718      for base_class in base_classes:
719        if issubclass(module_member, base_class):
720          subclasses.append(module_member)
721  return subclasses
722
723
724def find_subclass_in_module(base_class, module):
725  """Finds the single subclass of the given base class in the given module.
726
727  Args:
728    base_class: class, the base class to look for a subclass of in the module.
729    module: module, the module to look for the single subclass in.
730
731  Returns:
732    The single subclass of the given base class.
733
734  Raises:
735    ValueError: If the number of subclasses found was not exactly one.
736  """
737  subclasses = find_subclasses_in_module([base_class], module)
738  if len(subclasses) != 1:
739    raise ValueError(
740        'Expected 1 subclass of %s per module, found %s.'
741        % (base_class.__name__, [subclass.__name__ for subclass in subclasses])
742    )
743  return subclasses[0]
744