pakit_tests.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Automated bot for testing recipes.
"""
from __future__ import absolute_import, print_function
import argparse
from argparse import RawDescriptionHelpFormatter as RawDescriptionHelp
import glob
import mmap
import os
import re
import sys

__version__ = '0.1.0'
TEMPLATE = """\"\"\"
Do not edit this file.
Just rerun the program to regenerate tests.
\"\"\"
from __future__ import absolute_import
import os
import shutil
import subprocess
import tempfile


CONFIG_TEMPLATE = \"\"\"pakit:
  command:
    timeout: 120
  defaults:
    repo: stable
  log:
    enabled: true
    file: {0}
    level: debug
  paths:
    link: {1}
    prefix: {2}
    recipes: {3}
    source: {4}
  recipe:
    update_interval: 86400
    uris:
    - uri: test_recipes
{5}
\"\"\"


def write_config(tail):
    \"\"\"
    Write config for a test execution to path.

    Args:
        tail: Vaild yaml to affix to the end of CONFIG_TEMPLATE.

    Returns:
        Path to temp directory.
    \"\"\"
    root = tempfile.mkdtemp()
    recipe_d = os.path.join(root, 'recipe')
    os.mkdir(recipe_d)
    os.symlink('ROOT_RECS',
               os.path.join(recipe_d, 'test_recipes'))

    with open(os.path.join(root, 'pakit.yml'), 'w') as fout:
        fout.write(CONFIG_TEMPLATE.format(
            os.path.join(root, 'main.log'),
            os.path.join(root, 'link'),
            os.path.join(root, 'prefix'),
            recipe_d,
            os.path.join(root, 'source'),
            tail
        ))

    return root


def delete_it(path):
    \"\"\"
    File or folder, it is deleted.

    Args:
        path: path to a file or dir
    \"\"\"
    try:
        shutil.rmtree(path)
    except OSError:
        try:
            os.remove(path)
        except OSError:
            pass


class RecipeTest(object):
    def setup_method(self, method):
        recipe = type(self).__name__.replace('Test_', '').split('::')[0]
        repo = method.__name__.replace('test_', '')

        self.temp_d = write_config(recipe + ':\\n  repo: ' + repo)
        self.args = ['pakit', '--conf',
                     os.path.join(self.temp_d, 'pakit.yml'), 'install', recipe]
        self.new_env = os.environ.copy()
        new_path = os.environ['PATH'] + ':' + os.path.join(self.temp_d,
                                                           'link', 'bin')
        self.new_env.update({'PATH': new_path})

    def teardown_method(self, _):
        delete_it(self.temp_d)

"""


def create_args_parser():
    """
    Create the program argument parser.

    Returns:
        An argparse parser object.
    """
    prog_name = os.path.basename(os.path.dirname(sys.argv[0]))
    mesg = """
    This script will (re)generate tests for recipes.
    It will OVERWRITE existing tests.
    """
    mesg = mesg[0:-5]
    parser = argparse.ArgumentParser(prog=prog_name, description=mesg,
                                     formatter_class=RawDescriptionHelp)
    parser.add_argument('-v', '--version', action='version',
                        version='pakit_tests {0}'.format(__version__))
    parser.add_argument('recipes_root', help='the folder containing recipes')
    parser.add_argument('output', nargs='?', default='tests/test_recipes.py',
                        help="""relative path from recipes root to test file,
                        default: tests/test_recipes.py""")

    return parser


def extract_repo_names(text):
    """
    Given a string, extract all keys from the string.

    Returns:
        List of keys in the string.
    """
    matcher = re.compile(r'\'(\w+)\':')
    results = [matched.group(1) for matched in matcher.finditer(text)]
    return results


def extract_repo_block(text):
    """
    Given a string, extract ONLY the repos dictionary block.

    Returns:
        A string containing only required block.
    """
    return re.search(r'(self.repos\s*=\s*{.*?})', text, re.DOTALL).group(1)


def format_lines(recipes):
    """
    Transform the dictionary to lines to write.

    Returns:
        Lines to write to test file.
    """
    lines = []
    class_line = '\nclass Test_{0}(RecipeTest):'
    repo_line = """    def test_{0}(self):
        assert subprocess.call(self.args, cwd=self.temp_d,
                               env=self.new_env) == 0
"""

    for recipe in sorted(recipes):
        repo_name = recipes[recipe]
        lines.append(class_line.format(recipe))
        for repo_name in recipes[recipe]:
            lines.extend(repo_line.format(repo_name).split('\n'))

    return lines[0:-1]


def scan_recipes(recipe_d):
    """
    Scan the recipe directory and return relevant data.
    """
    data = {}
    matcher = re.compile(r'class\s+\S+\(Recipe\)')
    candidates = [fname for fname in glob.glob(os.path.join(recipe_d, '*.py'))]

    for fname in candidates:
        short_name = os.path.basename(fname)[0:-3]
        with open(fname, 'r+') as fin:
            text = mmap.mmap(fin.fileno(), 0)
        if matcher.search(text) is not None:
            data[short_name] = extract_repo_names(extract_repo_block(text))

    return data


def write_file(root, test_file):
    """
    Write the test file as requested.
    """
    try:
        os.makedirs(os.path.dirname(test_file))
    except OSError:
        pass

    text = TEMPLATE.replace('ROOT_RECS', root) + \
        '\n'.join(format_lines(scan_recipes(root)))
    with open(test_file, 'w') as fout:
        fout.write(text)


def main(argv=None):
    """
    The main entry point for this program.

    Args:
        argv: A list of program options, if None use sys.argv.
    """
    if argv is None:
        argv = sys.argv

    parser = create_args_parser()
    args = parser.parse_args(argv[1:])

    root = os.path.abspath(args.recipes_root)
    if os.path.isabs(args.output):
        test_file = os.path.join(root, args.output)
    else:
        test_file = os.path.join(root, args.output)

    print('Scanning recipes under: ' + root)
    print('Writing tests to: ' + test_file)
    write_file(root, test_file)


if __name__ == "__main__":
    main()  # pragma: no cover