# SPDX-License-Identifier: (GPL-2.0 OR Linux-OpenIB)
# Copyright (c) 2019 Mellanox Technologies, Inc. All rights reserved. See COPYING file
"""
Test module for pyverbs' mr module.
"""
from itertools import combinations as com
import random

from pyverbs.pyverbs_error import PyverbsRDMAError, PyverbsError
from pyverbs.tests.base import PyverbsTestCase
from pyverbs.base import PyverbsRDMAErrno
from pyverbs.mr import MR, MW, DMMR
import pyverbs.tests.utils as u
import pyverbs.device as d
from pyverbs.pd import PD
import pyverbs.enums as e

MAX_IO_LEN = 1048576


class MRTest(PyverbsTestCase):
    """
    Test various functionalities of the MR class.
    """
    def test_reg_mr(self):
        """ Test ibv_reg_mr() """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                flags = u.get_access_flags(ctx)
                for f in flags:
                    with MR(pd, u.get_mr_length(), f) as mr:
                        pass

    def test_dereg_mr(self):
        """ Test ibv_dereg_mr() """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                flags = u.get_access_flags(ctx)
                for f in flags:
                    with MR(pd, u.get_mr_length(), f) as mr:
                        mr.close()

    @staticmethod
    def test_reg_mr_bad_flow():
        """ Verify that trying to register a MR with None PD fails """
        try:
            # Use the simplest access flags necessary
            MR(None, random.randint(0, 10000), e.IBV_ACCESS_LOCAL_WRITE)
        except TypeError as te:
            assert 'expected pyverbs.pd.PD' in te.args[0]
            assert 'got NoneType' in te.args[0]
        else:
            raise PyverbsRDMAErrno('Created a MR with None PD')

    def test_dereg_mr_twice(self):
        """ Verify that explicit call to MR's close() doesn't fails """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                flags = u.get_access_flags(ctx)
                for f in flags:
                    with MR(pd, u.get_mr_length(), f) as mr:
                        # Pyverbs supports multiple destruction of objects,
                        # we are not expecting an exception here.
                        mr.close()
                        mr.close()

    def test_reg_mr_bad_flags(self):
        """ Verify that illegal flags combination fails as expected """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                for i in range(5):
                    flags = random.sample([e.IBV_ACCESS_REMOTE_WRITE,
                                           e.IBV_ACCESS_REMOTE_ATOMIC],
                                          random.randint(1, 2))
                    mr_flags = 0
                    for i in flags:
                        mr_flags += i.value
                    try:
                        MR(pd, u.get_mr_length(), mr_flags)
                    except PyverbsRDMAError as err:
                        assert 'Failed to register a MR' in err.args[0]
                    else:
                        raise PyverbsRDMAError('Registered a MR with illegal falgs')

    def test_write(self):
        """
        Test writing to MR's buffer
        """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                for i in range(10):
                    mr_len = u.get_mr_length()
                    flags = u.get_access_flags(ctx)
                    for f in flags:
                        with MR(pd, mr_len, f) as mr:
                            write_len = min(random.randint(1, MAX_IO_LEN),
                                            mr_len)
                            mr.write('a' * write_len, write_len)

    def test_read(self):
        """
        Test reading from MR's buffer
        """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                for i in range(10):
                    mr_len = u.get_mr_length()
                    flags = u.get_access_flags(ctx)
                    for f in flags:
                        with MR(pd, mr_len, f) as mr:
                            write_len = min(random.randint(1, MAX_IO_LEN),
                                            mr_len)
                            write_str = 'a' * write_len
                            mr.write(write_str, write_len)
                            read_len = random.randint(1, write_len)
                            offset = random.randint(0, write_len-read_len)
                            read_str = mr.read(read_len, offset).decode()
                            assert read_str in write_str

    def test_lkey(self):
        """
        Test reading lkey property
        """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                length = u.get_mr_length()
                flags = u.get_access_flags(ctx)
                for f in flags:
                    with MR(pd, length, f) as mr:
                        mr.lkey

    def test_rkey(self):
        """
        Test reading rkey property
        """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                length = u.get_mr_length()
                flags = u.get_access_flags(ctx)
                for f in flags:
                    with MR(pd, length, f) as mr:
                        mr.rkey

    def test_buffer(self):
        """
        Test reading buf property
        """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                length = u.get_mr_length()
                flags = u.get_access_flags(ctx)
                for f in flags:
                    with MR(pd, length, f) as mr:
                        mr.buf


class MWTest(PyverbsTestCase):
    """
    Test various functionalities of the MW class.
    """
    def test_reg_mw_type1(self):
        """ Test ibv_alloc_mw() """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                with MW(pd, e.IBV_MW_TYPE_1):
                    pass

    def test_reg_mw_type2(self):
        """ Test ibv_alloc_mw() """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                with MW(pd, e.IBV_MW_TYPE_2):
                    pass

    def test_dereg_mw_type1(self):
        """ Test ibv_dealloc_mw() """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                with MW(pd, e.IBV_MW_TYPE_1) as mw:
                    mw.close()

    def test_dereg_mw_type2(self):
        """ Test ibv_dealloc_mw() """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                with MW(pd, e.IBV_MW_TYPE_2) as mw:
                    mw.close()

    def test_reg_mw_wrong_type(self):
        """ Test ibv_alloc_mw() """
        for ctx, attr, attr_ex in self.devices:
            with PD(ctx) as pd:
                try:
                    mw_type = random.randint(3, 100)
                    MW(pd, mw_type)
                except PyverbsRDMAError:
                    pass
                else:
                    raise PyverbsError('Created a MW with type {t}'.\
                                       format(t=mw_type))


class DMMRTest(PyverbsTestCase):
    """
    Test various functionalities of the DMMR class.
    """
    def test_create_dm_mr(self):
        """
        Test ibv_reg_dm_mr
        """
        for ctx, attr, attr_ex in self.devices:
            if attr_ex.max_dm_size == 0:
                return
            with PD(ctx) as pd:
                for i in range(10):
                    dm_len = random.randrange(u.MIN_DM_SIZE, attr_ex.max_dm_size,
                                              u.DM_ALIGNMENT)
                    dm_attrs = u.get_dm_attrs(dm_len)
                    with d.DM(ctx, dm_attrs) as dm:
                        dm_mr_len = random.randint(1, dm_len)
                        dm_mr_offset = random.randint(0, (dm_len - dm_mr_len))
                        DMMR(pd, dm_mr_len, e.IBV_ACCESS_ZERO_BASED, dm=dm,
                             offset=dm_mr_offset)

    def test_destroy_dm_mr(self):
        """
        Test freeing of dm_mr
        """
        for ctx, attr, attr_ex in self.devices:
            if attr_ex.max_dm_size == 0:
                return
            with PD(ctx) as pd:
                for i in range(10):
                    dm_len = random.randrange(u.MIN_DM_SIZE, attr_ex.max_dm_size,
                                              u.DM_ALIGNMENT)
                    dm_attrs = u.get_dm_attrs(dm_len)
                    with d.DM(ctx, dm_attrs) as dm:
                        dm_mr_len = random.randint(1, dm_len)
                        dm_mr_offset = random.randint(0, (dm_len - dm_mr_len))
                        dm_mr = DMMR(pd, dm_mr_len, e.IBV_ACCESS_ZERO_BASED,
                                     dm=dm, offset=dm_mr_offset)
                        dm_mr.close()
