test.py

#!/usr/bin/env python

import glob
import os
import pickle
import re
import subprocess
import sys

import numexpr as ne
from iterfzf import iterfzf

CLASSNAME = re.compile(r'^class (?P<classname>\w+)\(')
TESTNAME = re.compile(r'\s+def (?P<testname>test_\w+)\(')

PARALLEL_THRESHOLD = 5
NUM_CORES = ne.detect_number_of_cores()
CWD_PATH = os.getcwd().lower()[1:].replace("/", ".")
LAST_COMMAND_PATH = f'/var/tmp/{CWD_PATH}_test_cmd'


def map_test_name_to_full_path():
    tests = {}

    for filename in glob.iglob('**/*.py', recursive=True):
        if not filename.endswith("tests.py"):
            continue

        modules = os.path.splitext(filename)[0].split('/')

        with open(filename) as f:
            class_name = None

            for line in f:
                match_data = CLASSNAME.match(line)
                if match_data:
                    class_name = match_data.group('classname')

                match_data = TESTNAME.match(line)
                if match_data:
                    test_name = match_data.group('testname')

                    app_name = modules[1]
                    example = test_name.split('test_')[1].replace('_', ' ')
                    test_case = class_name.split('Test')[0]
                    module_path = '.'.join(modules)

                    full_path = f'{module_path}.{class_name}.{test_name}'
                    index = f'{example} ({app_name}.{test_case})'
                    tests[index] = full_path

    return tests


def fuzzy_select_tests():
    test_mappings = map_test_name_to_full_path()

    _, selected_tests = iterfzf(
        test_mappings.keys(),
        multi=True,
        print_query=True,
    )

    if not selected_tests:
        exit(0)

    full_test_paths = [test_mappings[st] for st in selected_tests]

    return selected_tests, full_test_paths


def build_command(full_test_paths, parallelism_allowed=True, resetdb=False):
    command = 'docker-compose exec web python manage.py test'.split()

    if parallelism_allowed and len(full_test_paths) > PARALLEL_THRESHOLD:
        command.append(f'--parallel={NUM_CORES}')

    if not resetdb:
        command.append('--keepdb')

    return [*command, *full_test_paths]


def print_selections(selections):
    if len(selections) == 1:
        print('Running test:')
    else:
        print('Running tests:')

    for selection in selections:
        print(f'\t{selection}')

    print()


def error(message):
    print(f'Error: {message}. Re-select tests.', file=sys.stderr)


if __name__ == '__main__':
    selected_tests = []
    full_test_paths = []

    if '--last' in sys.argv:
        try:
            with open(LAST_COMMAND_PATH, 'rb') as f:
                selected_tests, full_test_paths = pickle.load(f)
        except FileNotFoundError:
            error('No previous test run found')
            selected_tests, full_test_paths = fuzzy_select_tests()
        except (KeyError, pickle.UnpicklingError):
            error('Could not read previous selection')
            selected_tests, full_test_paths = fuzzy_select_tests()
    else:
        selected_tests, full_test_paths = fuzzy_select_tests()

    with open(LAST_COMMAND_PATH, 'wb') as f:
        pickle.dump((selected_tests, full_test_paths), f)

    print_selections(selected_tests)

    command = build_command(
        full_test_paths,
        parallelism_allowed=('--no-parallel' not in sys.argv),
        resetdb=('--resetdb' in sys.argv),
    )

    if '--debug' in sys.argv:
        print(' '.join(command))

    subprocess.call(command)