1 files changed, 29 insertions, 21 deletions

diff --git a/run_tests b/run_tests
index d7144b3c..573dd446 100755
--- a/run_tests
+++ b/run_tests
@@ -1,5 +1,4 @@
-#!/usr/bin/env python
-# -*- coding:utf-8 -*-
+#!/usr/bin/env python3
 # Copyright 2019 The Android Open Source Project
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,37 +15,46 @@
 
 """Wrapper to run pytest with the right settings."""
 
-from __future__ import print_function
-
 import errno
 import os
+import shutil
 import subprocess
 import sys
 
 
-def run_pytest(cmd, argv):
-  """Run the unittests via |cmd|."""
-  try:
-    return subprocess.call([cmd] + argv)
-  except OSError as e:
-    if e.errno == errno.ENOENT:
-      print('%s: unable to run `%s`: %s' % (__file__, cmd, e), file=sys.stderr)
-      print('%s: Try installing pytest: sudo apt-get install python-pytest' %
-            (__file__,), file=sys.stderr)
-      return 127
-    else:
-      raise
+def find_pytest():
+  """Try to locate a good version of pytest."""
+  # If we're in a virtualenv, assume that it's provided the right pytest.
+  if 'VIRTUAL_ENV' in os.environ:
+    return 'pytest'
+
+  # Use the Python 3 version if available.
+  ret = shutil.which('pytest-3')
+  if ret:
+    return ret
+
+  # Hopefully this is a Python 3 version.
+  ret = shutil.which('pytest')
+  if ret:
+    return ret
+
+  print('%s: unable to find pytest.' % (__file__,), file=sys.stderr)
+  print('%s: Try installing: sudo apt-get install python-pytest' % (__file__,),
+        file=sys.stderr)
 
 
 def main(argv):
   """The main entry."""
   # Add the repo tree to PYTHONPATH as the tests expect to be able to import
   # modules directly.
-  topdir = os.path.dirname(os.path.realpath(__file__))
-  pythonpath = os.environ.get('PYTHONPATH', '')
-  os.environ['PYTHONPATH'] = '%s:%s' % (topdir, pythonpath)
-
-  return run_pytest('pytest', argv)
+  pythonpath = os.path.dirname(os.path.realpath(__file__))
+  oldpythonpath = os.environ.get('PYTHONPATH', None)
+  if oldpythonpath is not None:
+    pythonpath += os.pathsep + oldpythonpath
+  os.environ['PYTHONPATH'] = pythonpath
+
+  pytest = find_pytest()
+  return subprocess.run([pytest] + argv, check=False).returncode
 
 
 if __name__ == '__main__':