zulip/zerver/lib/test_runner.py

from __future__ import print_function
from django.test.runner import DiscoverRunner

from zerver.lib.cache import bounce_key_prefix_for_testing
from zerver.views.messages import get_sqlalchemy_connection

import os
import subprocess
import sys
import time
import traceback
import unittest

def slow(expected_run_time, slowness_reason):
    '''
    This is a decorate that annotates a test as being "known
    to be slow."  The decorator will set expected_run_time and slowness_reason
    as atributes of the function.  Other code can use this annotation
    as needed, e.g. to exclude these tests in "fast" mode.
    '''
    def decorator(f):
        f.expected_run_time = expected_run_time
        f.slowness_reason = slowness_reason
        return f

    return decorator

def is_known_slow_test(test_method):
    return hasattr(test_method, 'slowness_reason')

def full_test_name(test):
    test_module = test.__module__
    test_class = test.__class__.__name__
    test_method = test._testMethodName
    return '%s.%s.%s' % (test_module, test_class, test_method)

def get_test_method(test):
    return getattr(test, test._testMethodName)

def enforce_timely_test_completion(test_method, test_name, delay):
    if hasattr(test_method, 'expected_run_time'):
        # Allow for tests to run 50% slower than normal due
        # to random variations.
        max_delay = 1.5 * test_method.expected_run_time
    else:
        max_delay = 0.180 # seconds

    # Further adjustments for slow laptops:
    max_delay = max_delay * 3

    if delay > max_delay:
        print('Test is TOO slow: %s (%.3f s)' % (test_name, delay))

def fast_tests_only():
    return "FAST_TESTS_ONLY" in os.environ

def run_test(test):
    failed = False
    test_method = get_test_method(test)

    if fast_tests_only() and is_known_slow_test(test_method):
        return failed

    test_name = full_test_name(test)

    bounce_key_prefix_for_testing(test_name)

    print('Running', test_name)
    if not hasattr(test, "_pre_setup"):
        # test_name is likely of the form unittest.loader.ModuleImportFailure.zerver.tests.test_upload
        import_failure_prefix = 'unittest.loader.ModuleImportFailure.'
        if test_name.startswith(import_failure_prefix):
            actual_test_name = test_name[len(import_failure_prefix):]
            print()
            print("Actual test to be run is %s, but import failed." % (actual_test_name,))
            print("Importing test module directly to generate clearer traceback:")
            try:
                command = ["python2.7", "-c", "import %s" % (actual_test_name,)]
                print("Import test command: `%s`" % (' '.join(command),))
                subprocess.check_call(command)
            except subprocess.CalledProcessError:
                print("If that traceback is confusing, try doing the import inside `./manage.py shell`")
                print()
                return True
            print("Import unexpectedly succeeded!  Something is wrong")
            return True
        else:
            print("Test doesn't have _pre_setup; something is wrong.")
            print("Here's a debugger. Good luck!")
            import pdb; pdb.set_trace()
    test._pre_setup()

    start_time = time.time()

    test.setUp()
    try:
        test_method()
    except unittest.SkipTest:
        pass
    except Exception:
        failed = True
        traceback.print_exc()

    test.tearDown()

    delay = time.time() - start_time
    enforce_timely_test_completion(test_method, test_name, delay)

    test._post_teardown()
    return failed

class Runner(DiscoverRunner):
    def __init__(self, *args, **kwargs):
        DiscoverRunner.__init__(self, *args, **kwargs)

    def run_suite(self, suite, fatal_errors=None):
        failed = False
        for test in suite:
            if run_test(test):
                failed = True
                if fatal_errors:
                    return failed
        return failed

    def run_tests(self, test_labels, extra_tests=None, **kwargs):
        self.setup_test_environment()
        try:
            suite = self.build_suite(test_labels, extra_tests)
        except AttributeError:
            traceback.print_exc()
            print()
            print("  This is often caused by a test module/class/function that doesn't exist or ")
            print("  import properly. You can usually debug in a `manage.py shell` via e.g. ")
            print("    import zerver.tests.test_messages")
            print("    from zerver.tests.test_messages import StreamMessagesTest")
            print("    StreamMessagesTest.test_message_to_stream")
            print()
            sys.exit(1)
        # We have to do the next line to avoid flaky scenarios where we
        # run a single test and getting an SA connection causes data from
        # a Django connection to be rolled back mid-test.
        get_sqlalchemy_connection()
        failed = self.run_suite(suite, fatal_errors=kwargs.get('fatal_errors'))
        self.teardown_test_environment()
        return failed
        print()