testmon.sh

testmon.sh

#!/usr/bin/env bash

download_database=false
upload_database=false
collect_only=false
run_testmon=false
incremental_collect=false
export_impacted_tests=false
input_sha=$CIRCLE_SHA1
S3_BUCKET="YOUR_S3_BUCKET"
S3_PREFIX="S3_BUCKET_PATH"
pytest_args=""

for arg in "$@"; do
    case $arg in
    --download-database)
        download_database=true
        shift
        ;;
    --upload-database)
        upload_database=true
        shift
        ;;
    --collect-only)
        collect_only=true
        shift
        ;;
    --incremental)
        incremental_collect=true
        shift
        ;;
    --run)
        run_testmon=true
        shift
        ;;
    --export-impacted-tests)
        export_impacted_tests=true
        shift
        ;;
    *)
        # Handle unknown options
        pytest_args="$pytest_args $arg"
        ;;
    esac
done

get_commit_list() {
    # Get the list of commit SHAs from the S3 bucket
    commit_list=$(aws s3api list-objects-v2 --bucket "$S3_BUCKET" --prefix "$S3_PREFIX" --query 'Contents | sort_by(@, &LastModified) | reverse(@) | [0:200].[Key]' --output text | grep -v "$input_sha" | grep ".testmondata$" | sed "s/$S3_PREFIX\///" | sed 's/\/.testmondata$//')
    echo "$commit_list"
}

# Function to find the closest parent commit
find_closest_parent() {
    local input_sha=$1
    local commit_list=$2
    local closest_parent=""

    for commit in $commit_list; do
        # Check if the commit is an ancestor of the input commit
        if git merge-base --is-ancestor "$commit" "$input_sha"; then
            # If it's the first ancestor found or it's a more recent ancestor, update the closest_parent variable
            if [[ -z $closest_parent ]] || git merge-base --is-ancestor "$closest_parent" "$commit"; then
                closest_parent=$commit
            fi
        fi
    done

    echo "$closest_parent"
}

download_testmondata() {
    # Find the closest parent commit and print it
    commit_list=$(get_commit_list)

    commits_len=$(echo "$commit_list" | wc -l)
    echo "📋 Last $commits_len commit list: $commit_list"

    closest_parent=$(find_closest_parent "$input_sha" "$commit_list")

    if [[ -z "$closest_parent" ]]; then
        echo "ℹ️  No closest parent found. Branch ${CIRCLE_BRANCH} is too old"
        exit 1
    else
        echo "ℹ️  Closest parent commit: '$closest_parent'"
    fi

    aws s3 cp --recursive "s3://$S3_BUCKET/$S3_PREFIX/${closest_parent}/" .testmondata-metadata
    mv .testmondata-metadata/.testmondata .testmondata

    # print size of the testmondata file
    echo "🗄️ Database file info"
    ls -lah .testmondata
}

export_impacted_tests() {
    download_testmondata
    mkdir -p reports
    pytest -p no:sugar -n1 --testmon-nocollect \
        --collect-only -qq | awk -F: '{print $1}' | grep -v "generated xml file" |
        grep -v "^$" | uniq >reports/pytest_impacted_tests.txt
    TOTAL_IMPACTED_TESTS=$(wc -l reports/pytest_impacted_tests.txt | awk '{print $1}')
    echo "Total number of tests impacted: $TOTAL_IMPACTED_TESTS"
}

run_impacted_tests() {
    TEST_FILES=""
    if [[ -f reports/pytest_impacted_tests.txt ]]; then
        TEST_FILES=$(tr <reports/pytest_impacted_tests.txt ' ' '\n' | sort | uniq)
    else
        TEST_FILES=$(circleci tests glob "**/{test,tests}/**/test_*.py" | sed 's/\S*venv.*//g' | sed 's/\S*__init__.py//g' | tr ' ' '\n' | sort | uniq)
    fi

    if [[ -z "$TEST_FILES" ]]; then
        echo "ℹ️  No test files to run. Exiting"
        exit 0
    fi

    echo "$TEST_FILES" >circleci_test_files.txt

    # run the tests
    circleci <circleci_test_files.txt tests run --command="xargs pytest -p no:sugar $pytest_args" --split-by=timings

    pytest_exit_code=$?

    # ignore only exit code 5 (no tests collected)
    if [ $pytest_exit_code -ne 5 ]; then
        echo "pytest exited with $pytest_exit_code"
        exit $pytest_exit_code
    fi
}

if $run_testmon; then
    run_impacted_tests
elif $download_database; then
    echo "⤵️  Downloading testmondata"
    download_testmondata
elif $upload_database; then
    mv .testmondata .testmondata-metadata/.testmondata
    echo "⤴️  Uploading testmondata"
    aws s3 cp .testmondata-metadata "s3://$S3_BUCKET/$S3_PREFIX/$input_sha" --recursive
    echo "✅ Uploaded testmondata"
elif $collect_only; then
    if $incremental_collect; then
        download_testmondata
    fi
    pytest -p no:sugar -n16 --reruns 4 --testmon
    pytest_exit_code=$?

    # Proceed with the upload of the testmondata file even if the tests fail
    # https://github.com/pytest-dev/pytest/blob/main/doc/en/reference/exit-codes.rst
    if [ $pytest_exit_code -eq 1 ]; then
        echo "ℹ️  pytest exited with 1"
        exit 0
    fi
    exit $pytest_exit_code
elif $export_impacted_tests; then
    download_testmondata
    export_impacted_tests
else
    echo "❗️ Please specify --run, --download-database, --upload-database, --collect-only"
    echo "ℹ️  --run: run the tests using testmon (--testmon-nocollect --testmon-forceselect)"
    echo "ℹ️  --download-database: download the testmondata file from S3"
    echo "ℹ️  --upload-database: upload the testmondata file to S3"
    echo "ℹ️  --collect-only: collect the testmondata file by running the tests"
    echo "ℹ️  --export-impacted-tests: export the list of impacted tests to reports/pytest_impacted_tests.txt"
    exit 1
fi