#!/bin/bash

# Run python unit tests

set -e

DEB_HOST_ARCH=$(dpkg-architecture -qDEB_HOST_ARCH)
DEB_HOST_ARCH_BITS=$(dpkg-architecture -qDEB_HOST_ARCH_BITS)

PYVER=$(py3versions -sv)

# Tests real build by default.
# Access the complex build by setting env variable USE_DOLFINX_COMPLEX

if [ "x${USE_DOLFINX_COMPLEX}" != "x" ]; then
    TESTDIR_SUFFIX="-complex"
    TEST_DESCRIPTION="for complex build "
    export PETSC_DIR=/usr/lib/petscdir/petsc-complex
    # process ufl forms for complex test: second argument 1 indicates complex
    UFL_TYPE=1
else
    TESTDIR_SUFFIX="-real"
    export PETSC_DIR=/usr/lib/petsc
fi
export SLEPC_DIR=`grep wPETSC_DIR ${PETSC_DIR}/lib/petsc/conf/petscvariables | awk '{print $3}' | sed "s/petsc/slepc/g"`

# MPI tests are set up to run on no more than 3 processes
NPROC=$( nproc )
if [[ $NPROC > 2 ]]; then
  N_MPI=3
else
  N_MPI=2
fi
# and only 2 processes for unit tests
N_MPI_UNITTEST=2

export OMPI_MCA_plm_rsh_agent=/bin/false
export OMPI_MCA_rmaps_base_oversubscribe=1
export OMPI_MCA_btl_base_warn_component_unused=0

# dolfinx needs Adios2 with MPI support
export ADIOS2_ALWAYS_USE_MPI=1

echo "== running python tests ${TEST_DESCRIPTION}=="

# help CI tests run a little faster
export DOLFINX_JIT_CFLAGS="-g0 -O0"

# numba is not available on all arches
if ! dpkg-query -s python3-numba >/dev/null 2>&1; then
    NUMBA_TESTS="custom_jit_kernels custom_assembler expression_evaluation"
    EXCLUDE_NUMBA=
    for nt in $NUMBA_TESTS; do
	EXCLUDE_NUMBA="$EXCLUDE_NUMBA --ignore-glob=*${nt}*"
    done
fi

declare -a SKIP_TEST_LIST

# GJK (test_cube_distance in test_gjk) fails to converge on many (but not all) arches
case " armel armhf mips64el sparc64 " in
  *\ ${DEB_HOST_ARCH}\ *)  SKIP_TEST_LIST=(${SKIP_TEST_LIST[@]} test_cube_distance);;
esac

# test_nonlinear_pde does not converge on some arches
# "Newton solver did not converge because maximum number of iterations reached"
case " mips64el ppc64 " in
  *\ ${DEB_HOST_ARCH}\ *)  SKIP_TEST_LIST=(${SKIP_TEST_LIST[@]} test_nonlinear_pde);;
esac

# riscv64 times out when running tests
# so skip the slowest of them, which is most of them
if [ "x$DEB_HOST_ARCH" = "xriscv64" ]; then
    # run verbose to monitor failing tests
    VERBOSE_TESTS="-v"
    SKIP_TEST_LIST=(${SKIP_TEST_LIST[@]} test_symmetry test_curl_curl_eigenvalue \
	test_higher_order_function test_eigen_assembly test_xdmf_function \
	test_assembler test_fem_pipeline test_dof_permuting test_dofmap test_element_integrals\
	test_assemble_submesh test_interpolation test_petsc_discrete_operators \
	test_mesh test_higher_order_mesh)
fi

TEST_KEYWORDS=""
for t in ${SKIP_TEST_LIST[@]}; do
    if [ "x${TEST_KEYWORDS}" = "x" ]; then
        TEST_KEYWORDS=$t
    else
        TEST_KEYWORDS="${TEST_KEYWORDS} or $t"
    fi
done
if [ "x${TEST_KEYWORDS}" != "x" ]; then
    TEST_KEYWORDS="not ( ${TEST_KEYWORDS} )"
fi
echo "skipping tests with SKIP_TEST_LIST=${SKIP_TEST_LIST[@]}"

for pyver in $PYVER; do
  echo "=== python $pyver unit test (serial) ${TEST_DESCRIPTION}==="
  python$pyver -m pytest --durations=20 $VERBOSE_TESTS $EXCLUDE_NUMBA -k "${TEST_KEYWORDS}" python/test/unit/
done

# run MPI tests
declare -a MPI_SKIP_TEST_LIST

if [ "x${DEB_HOST_ARCH}" = "xs390x" ]; then
  MPI_SKIP_TEST_LIST=(${MPI_SKIP_TEST_LIST[@]} test_fem_pipeline)
fi

MPI_TEST_KEYWORDS=""
for t in ${SKIP_TEST_LIST[@]} ${MPI_SKIP_TEST_LIST[@]}; do
    if [ "x${MPI_TEST_KEYWORDS}" = "x" ]; then
        MPI_TEST_KEYWORDS=$t
    else
        MPI_TEST_KEYWORDS="${MPI_TEST_KEYWORDS} or $t"
    fi
done
if [ "x${MPI_TEST_KEYWORDS}" != "x" ]; then
    MPI_TEST_KEYWORDS="not ( ${MPI_TEST_KEYWORDS} )"
fi
echo "also skipping MPI tests with MPI_SKIP_TEST_LIST=${MPI_SKIP_TEST_LIST[@]}"

if [ "x${DEB_HOST_ARCH_BITS}" = "x32" ]; then
   echo "=== python unit tests (MPI) skipped on 32-bit systems (see Bug#995599) ${TEST_DESCRIPTION}==="
else
  for pyver in $PYVER; do
    echo "=== python $pyver unit test (MPI using ${N_MPI_UNITTEST} processors out of $NPROC) ${TEST_DESCRIPTION}==="
    mpirun -n ${N_MPI_UNITTEST} python$pyver -m pytest --durations=20 $VERBOSE_TESTS $EXCLUDE_NUMBA -k "${MPI_TEST_KEYWORDS}" python/test/unit/ --color=no
  done
fi
