import tempfile
from h5py import File
import h5py

mpi = h5py.get_config().mpi

if not mpi:
    print("Parallel HDF5 is required for MPIO driver tests")
    exit(1)

def mktemp_mpi(comm=None, suffix='.hdf5', prefix='', dir=None):
    if comm is None:
        from mpi4py import MPI
        comm = MPI.COMM_WORLD
    fname = None
    if comm.Get_rank() == 0:
        fname = tempfile.mktemp(suffix, prefix, dir)
    fname = comm.bcast(fname, 0)
    return fname


def test_mpio():
    """ MPIO driver and options """
    from mpi4py import MPI

    comm=MPI.COMM_WORLD
    fname = mktemp_mpi(comm)
    with File(fname, 'w', driver='mpio', comm=comm) as f:
        assert f
        assert f.driver=='mpio'

def test_mpi_atomic():
    """ Enable atomic mode for MPIO driver """
    from mpi4py import MPI

    comm=MPI.COMM_WORLD
    fname = mktemp_mpi(comm)
    with File(fname, 'w', driver='mpio', comm=comm) as f:
        assert not f.atomic
        f.atomic = True
        assert f.atomic


def test_close_multiple_mpio_driver():
    """ MPIO driver and options """
    from mpi4py import MPI

    comm=MPI.COMM_WORLD
    fname = mktemp_mpi(comm)
    f = File(fname, 'w', driver='mpio', comm=comm)
    f.create_group("test")
    f.close()
    f.close()


test_mpio()
test_mpi_atomic()
test_close_multiple_mpio_driver()
