#!/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

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

TESTS_SKIPPED="test_RT_N1curl_simplex"
MULTIPLE_TESTS_SKIPPED="test_assembly_solve_taylor_hood test_cube_distance test_read_write_p2_mesh gmsh"

for test in $MULTIPLE_TESTS_SKIPPED $TESTS_SKIPPED_32BIT; do
  TESTS_SKIPPED="$TESTS_SKIPPED or $test"
done
TEST_KEYWORDS="not (${TESTS_SKIPPED})"

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

# run MPI tests

if [ "x${DEB_HOST_ARCH}" = "xs390x" ]; then
  OTHER_MPI_TESTS_SKIPPED="test_fem_pipeline"
fi

MPI_TESTS_SKIPPED=$TESTS_SKIPPED
for mpitest in $OTHER_MPI_TESTS_SKIPPED; do
  MPI_TESTS_SKIPPED="$MPI_TESTS_SKIPPED or $mpitest"
done
MPI_TEST_KEYWORDS="not ( $MPI_TESTS_SKIPPED )"

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 $EXCLUDE_NUMBA -k "${MPI_TEST_KEYWORDS}" python/test/unit/ --color=no
  done
fi
