Restore running of benchmarks in test_snippets.py

Daniel Watkins · Daniel Watkins · commit c857d7574865 · 2018-08-04T12:51:12.000-04:00
This distinguishes them from regular snippet tests, and runs them
without trace logging (which was causing major slowdowns previously).
diff --git a/tests/test_snippets.py b/tests/test_snippets.py
@@ -10,13 +10,23 @@
 import logging
 import subprocess
 import contextlib
+import enum
 
 import compile_code
 
 
+class _TestType(enum.Enum):
+    functional = 1
+    benchmark = 2
+
+
 logger = logging.getLogger('tests')
 ROOT_DIR = '..'
-TEST_DIR = os.path.abspath(os.path.join(ROOT_DIR, 'tests', 'snippets'))
+TEST_ROOT = os.path.abspath(os.path.join(ROOT_DIR, 'tests'))
+TEST_DIRS = {
+    _TestType.functional: os.path.join(TEST_ROOT, 'snippets'),
+    _TestType.benchmark: os.path.join(TEST_ROOT, 'benchmarks'),
+}
 CPYTHON_RUNNER_DIR = os.path.abspath(os.path.join(ROOT_DIR, 'py_code_object'))
 RUSTPYTHON_RUNNER_DIR = os.path.abspath(os.path.join(ROOT_DIR))
 
@@ -29,14 +39,14 @@ def pushd(path):
     os.chdir(old_dir)
 
 
-def perform_test(filename, method):
+def perform_test(filename, method, test_type):
     logger.info('Running %s via %s', filename, method)
     if method == 'cpython':
         run_via_cpython(filename)
     elif method == 'cpython_bytecode':
-        run_via_cpython_bytecode(filename)
+        run_via_cpython_bytecode(filename, test_type)
     elif method == 'rustpython':
-        run_via_rustpython(filename)
+        run_via_rustpython(filename, test_type)
     else:
         raise NotImplementedError(method)
 
@@ -46,37 +56,39 @@ def run_via_cpython(filename):
     subprocess.check_call([sys.executable, filename])
 
 
-def run_via_cpython_bytecode(filename):
+def run_via_cpython_bytecode(filename, test_type):
     # Step1: Create bytecode file:
     bytecode_filename = filename + '.bytecode'
     with open(bytecode_filename, 'w') as f:
         compile_code.compile_to_bytecode(filename, out_file=f)
 
     # Step2: run cpython bytecode:
     env = os.environ.copy()
-    env['RUST_LOG'] = 'debug,cargo=error,jobserver=error'
+    log_level = 'info' if test_type == _TestType.benchmark else 'debug'
+    env['RUST_LOG'] = '{},cargo=error,jobserver=error'.format(log_level)
     env['RUST_BACKTRACE'] = '1'
     with pushd(CPYTHON_RUNNER_DIR):
         subprocess.check_call(['cargo', 'run', bytecode_filename], env=env)
 
 
-def run_via_rustpython(filename):
+def run_via_rustpython(filename, test_type):
     env = os.environ.copy()
-    env['RUST_LOG'] = 'trace,cargo=error,jobserver=error'
+    log_level = 'info' if test_type == _TestType.benchmark else 'trace'
+    env['RUST_LOG'] = '{},cargo=error,jobserver=error'.format(log_level)
     env['RUST_BACKTRACE'] = '1'
     with pushd(RUSTPYTHON_RUNNER_DIR):
         subprocess.check_call(['cargo', 'run', '--release', filename], env=env)
 
 
-def create_test_function(cls, filename, method):
+def create_test_function(cls, filename, method, test_type):
     """ Create a test function for a single snippet """
     core_test_directory, snippet_filename = os.path.split(filename)
     test_function_name = 'test_{}_'.format(method) \
         + os.path.splitext(snippet_filename)[0] \
         .replace('.', '_').replace('-', '_')
 
     def test_function(self):
-        perform_test(filename, method)
+        perform_test(filename, method, test_type)
 
     if hasattr(cls, test_function_name):
         raise ValueError('Duplicate test case {}'.format(test_function_name))
@@ -86,21 +98,21 @@ def test_function(self):
 def populate(method):
     def wrapper(cls):
         """ Decorator function which can populate a unittest.TestCase class """
-        for filename in get_test_files():
-            create_test_function(cls, filename, method)
+        for test_type, filename in get_test_files():
+            create_test_function(cls, filename, method, test_type)
         return cls
     return wrapper
 
 
 def get_test_files():
     """ Retrieve test files """
-    for filepath in sorted(glob.iglob(os.path.join(
-            TEST_DIR, '*.py'))):
-        filename = os.path.split(filepath)[1]
-        if filename.startswith('xfail_'):
-            continue
+    for test_type, test_dir in TEST_DIRS.items():
+        for filepath in sorted(glob.iglob(os.path.join(test_dir, '*.py'))):
+            filename = os.path.split(filepath)[1]
+            if filename.startswith('xfail_'):
+                continue
 
-        yield os.path.abspath(filepath)
+            yield test_type, os.path.abspath(filepath)
 
 
 @populate('cpython')
