Автоматическая валидация кода python-тестов через ast-модуль

Привет всем!

Хотелось бы поделиться, на мой взгляд, интересной и полезной темой, как кастомные валидаторы кода для автоматической проверки тестов.

Зачем это нужно?

Думаю, что многие python-автоматизаторы знают и используют популярную библиотеку flake8, которая имплементирует pep8-верификацию, или же pylint-библиотеку. Синтаксические анализаторы делают полезную службу тем, что помогают исправлять синтаксические ошибки (плюс следовать общему стилю кодирования) на стадии написания кода, до того как код уедет в CI. Это быстро, это удобно, это дешево. Бывает обидно поймать опечатку в тесте на CI, где деплой тестируемого облака занимает не менее часа.

В любом проекте всегда возникают правила (conventions), которым следует команда (гласно или негласно, лучше если гласно, общепринято и задокументировано). Единообразие кода в большой и распределенной команде является значимой темой. При этом классно иметь автоматические средства валидации принятых правил.

Для нашего проекта одним из правил является то, что в коде тестов мы вызываем только стэпы (и py.test-фикстуры), для того чтобы код читался максимально понятно (самодокументируемый код). Часто можно встретить тесты, в которых выполняются различные низкоуровневые операции, делающие код теста плохочитаемым, и которые на деле можно вынести в общие переиспользуемые методы. Чтобы не отвлекаться на этот момент на ревью, было решено сделать автоматическую проверку консистентности тестов перед запуском.

Как это реализовано?

Т.к. мы используем py.test для написания и запуска автотестов, steps-checker было решено делать как py.test-плагин. Вот его код:

def pytest_collection_modifyitems(config, items):
    """Hook to detect forbidden calls inside test."""
    if config.option.disable_steps_checker:
        config.warn('P1', 'Permitted calls checker is disabled!')
        return

    errors = []
    for item in items:
        permitted_calls = PERMITTED_CALLS + STEPS + item.funcargnames

        ast_root = ast.parse(_get_source(item.function))
        for call_name in _get_call_names(ast_root):
            if call_name not in permitted_calls:

                error = ("Calling {!r} isn't allowed".format(call_name) +
                         _get_func_location(item.function) + DOC_LINK)
                errors.append(error)

    if errors:
        pytest.exit("Only steps and fixtures must be called in test!\n" +
                    '\n'.join(errors))

Мы использовали модуль ast, который умеет строить синтактическое дерево кода, получая на вход текст кода. К сожалению, модуль плохо документирован. Но как оказалось ничего сложного в нем, нотация весьма понятна и разобраться при желании удалось за 30 мин.

Мы рекурсивно обходим дерево кода теста и ищем все вызовы (ast.Call) и затем проверяем, что среди найденных вызовов отсутствуют неразрешенные. Вот код, который имплементирует рекурсивный обход дерева.

def _get_call_names(node):
    """Get called function names inside function."""
    call_names = set()
    call_nodes = _get_ast_nodes(node, ast.Call)

    for call_node in call_nodes:
        if hasattr(call_node.func, 'attr'):
            call_name = call_node.func.attr

        else:
            call_name = call_node.func.id
        call_names.add(call_name)

    return call_names


def _get_ast_nodes(node, node_type, bucket=None):
    """Get ast nodes with specifed ast type.

    Recursive traversal of ast nodes tree to retrieve nodes by defined type.
    """
    bucket = [] if bucket is None else bucket

    if isinstance(node, node_type):
        bucket.append(node)

    for attr in _get_ast_attrs(node):
        if not utils.is_iterable(attr):
            attr = [attr]

        for elem in attr:
            if isinstance(elem, ast.AST):
                if isinstance(elem, node_type):
                    bucket.append(elem)
                bucket = _get_ast_nodes(elem, node_type, bucket)
    return bucket


def _get_ast_attrs(node):
    """Get attributes of ast node.

    Some attributes are skipped because they are not related to function body
    and just create side effect.
    """
    skip_list = ()
    if isinstance(node, ast.FunctionDef):
        skip_list = ('args', 'decorator_list')

    attrs = []
    for name in dir(node):
        if not name.startswith('_') and name not in skip_list:
            attrs.append(getattr(node, name))
    return attrs

На CI у нас настроена джоба, которая выполняет команду коллекции тестов py.test ${PROJECT_DIR} --collect-only, вовлекающую наш плагин. Команду также можно (и нужно!) выполнять и локально на машине автоматизатора.

Подробнее код можно посмотреть в stepler/steps_checker.py at master · Mirantis/stepler · GitHub.