#!/usr/bin/env python
# Don't run tests from the root repo dir.
# We want to ensure we're importing from the installed
# binary package not from the CWD.

import argparse
import os
from contextlib import contextmanager
from subprocess import check_call

_dname = os.path.dirname

REPO_ROOT = _dname(_dname(_dname(os.path.abspath(__file__))))
PACKAGE = "s3transfer"


@contextmanager
def cd(path):
    """Change directory while inside context manager."""
    cwd = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(cwd)


def run(command, env=None):
    return check_call(command, shell=True, env=env)


def process_args(args):
    runner = args.test_runner
    test_args = ""
    if args.with_cov:
        test_args += f"--cov={PACKAGE} --cov-report xml "
    dirs = " ".join(args.test_dirs)

    return runner, test_args, dirs


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "test_dirs",
        default=["unit/", "functional/"],
        nargs="*",
        help="One or more directories containing tests.",
    )
    parser.add_argument(
        "-r",
        "--test-runner",
        default="pytest",
        help="Test runner to execute tests. Defaults to pytest.",
    )
    parser.add_argument(
        "-c",
        "--with-cov",
        default=False,
        action="store_true",
        help="Run default test-runner with code coverage enabled.",
    )
    raw_args = parser.parse_args()
    test_runner, test_args, test_dirs = process_args(raw_args)

    with cd(os.path.join(REPO_ROOT, "tests")):
        cmd = f"{test_runner} {test_args}{test_dirs}"
        print(f"Running {cmd}...")
        run(cmd)

        # Run the serial implementation of s3transfer
        os.environ["USE_SERIAL_EXECUTOR"] = "True"
        print(f"Running serial execution for {cmd}...")
        run(cmd, env=os.environ)