/* SPDX-License-Identifier: BSD-2-Clause OR GPL-2.0-only */
/* SPDX-FileCopyrightText: Copyright Amazon.com, Inc. or its affiliates. All rights reserved. */

#include "efa_unit_tests.h"

static void test_efa_mr_impl(struct efa_domain *efa_domain, struct fid_mr *mr,
			int mr_reg_count, int mr_reg_size)
{
	assert_int_equal(ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct), (int64_t)mr_reg_count);
	assert_int_equal(ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz), (int64_t)mr_reg_size);
}

static void test_efa_rdm_mr_impl(struct efa_domain *efa_domain, struct fid_mr *mr,
				int mr_reg_count, int mr_reg_size, bool gdrcopy_flag)
{
	struct efa_rdm_mr *efa_rdm_mr;

	test_efa_mr_impl(efa_domain, mr, mr_reg_count, mr_reg_size);

	if (mr) {
		efa_rdm_mr = container_of(mr, struct efa_rdm_mr, efa_mr.mr_fid);

		/* Test RDM-specific MR map insertion status */
		assert_true(efa_rdm_mr->inserted_to_mr_map);

		/* Test SHM MR creation based on domain availability */
		if (efa_domain->shm_domain) {
			assert_non_null(efa_rdm_mr->shm_mr);
		} else {
			assert_null(efa_rdm_mr->shm_mr);
		}

		/* Test needs_sync flag for CUDA memory */
		if (efa_rdm_mr->efa_mr.iface == FI_HMEM_CUDA) {
			assert_true(efa_rdm_mr->needs_sync);
		} else {
			assert_false(efa_rdm_mr->needs_sync);
		}

		if (cuda_is_gdrcopy_enabled()) {
			if (gdrcopy_flag) {
				assert_true(efa_rdm_mr->flags &
					    OFI_HMEM_DATA_DEV_REG_HANDLE);
				assert_non_null(efa_rdm_mr->hmem_data);
			} else {
				assert_false(efa_rdm_mr->flags &
					     OFI_HMEM_DATA_DEV_REG_HANDLE);
				assert_null(efa_rdm_mr->hmem_data);
			}
		}
	}
}

void test_efa_rdm_mr_reg_host_memory(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	int baseline_ct, baseline_sz;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	/* fi_endpoint calls ofi_bufpool_grow, which registers mr */
	baseline_ct = ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct);
	baseline_sz = ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz);

	buf = malloc(mr_size);
	assert_non_null(buf);

	assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
				   FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL),
			 0);

	/* No GDRCopy registration for host memory */
	test_efa_rdm_mr_impl(efa_domain, mr, baseline_ct + 1, baseline_sz + mr_size, false);

	assert_int_equal(fi_close(&mr->fid), 0);
	test_efa_rdm_mr_impl(efa_domain, NULL, baseline_ct, baseline_sz, false);
	free(buf);
}

void test_efa_rdm_mr_reg_host_memory_no_mr_local(void **state)
{
	struct efa_resource *resource = *state;
	struct fi_info *hints;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	int baseline_ct, baseline_sz;

	hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode &= ~FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						 FI_VERSION(2, 0), hints, true, true);
	fi_freeinfo(hints);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	/* fi_endpoint calls ofi_bufpool_grow, which registers mr */
	baseline_ct = ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct);
	baseline_sz = ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz);

	buf = malloc(mr_size);
	assert_non_null(buf);

	assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
				   FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL),
			 0);

	/* No GDRCopy registration for host memory */
	test_efa_rdm_mr_impl(efa_domain, mr, baseline_ct + 1, baseline_sz + mr_size, false);

	assert_int_equal(fi_close(&mr->fid), 0);
	test_efa_rdm_mr_impl(efa_domain, NULL, baseline_ct, baseline_sz, false);
	free(buf);
}

void test_efa_rdm_mr_reg_host_memory_overlapping_buffers(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	size_t mr_size = 1024;
	void *buf;
	struct fid_mr *mr_1 = NULL, *mr_2 = NULL;
	int baseline_ct, baseline_sz;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	/* fi_endpoint calls ofi_bufpool_grow, which registers mr */
	baseline_ct = ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct);
	baseline_sz = ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz);

	buf = malloc(mr_size);
	assert_non_null(buf);

	assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
				   FI_SEND | FI_RECV, 0, 0, 0, &mr_1, NULL),
			 0);

	/* No GDRCopy registration for host memory */
	test_efa_rdm_mr_impl(efa_domain, mr_1, baseline_ct + 1, baseline_sz + mr_size, false);

	assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
				   FI_SEND | FI_RECV, 0, 0, 0, &mr_2, NULL),
			 0);

	test_efa_rdm_mr_impl(efa_domain, mr_2, baseline_ct + 2, baseline_sz + mr_size * 2, false);

	assert_int_equal(fi_close(&mr_1->fid), 0);
	test_efa_rdm_mr_impl(efa_domain, mr_2, baseline_ct + 1, baseline_sz + mr_size, false);

	assert_int_equal(fi_close(&mr_2->fid), 0);
	test_efa_rdm_mr_impl(efa_domain, NULL, baseline_ct, baseline_sz, false);

	free(buf);
}

#if HAVE_CUDA
void test_efa_rdm_mr_reg_cuda_memory(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	struct fi_mr_attr mr_reg_attr = { 0 };
	struct iovec iovec;
	int err, baseline_ct, baseline_sz;

	if (hmem_ops[FI_HMEM_CUDA].initialized &&
	    g_efa_hmem_info[FI_HMEM_CUDA].p2p_supported_by_device) {
		resource->hints = efa_unit_test_alloc_hints_hmem(
			FI_EP_RDM, EFA_FABRIC_NAME);
		efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
							    FI_VERSION(2, 0),
							    resource->hints,
							    true, true);

		efa_domain = container_of(resource->domain, struct efa_domain,
					  util_domain.domain_fid);
		/* fi_endpoint calls ofi_bufpool_grow, which registers mr */
		baseline_ct = ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct);
		baseline_sz = ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz);

		err = ofi_cudaMalloc(&buf, mr_size);
		assert_int_equal(err, 0);
		assert_non_null(buf);

		mr_reg_attr.access = FI_SEND | FI_RECV;
		mr_reg_attr.iface = FI_HMEM_CUDA;
		iovec.iov_base = buf;
		iovec.iov_len = mr_size;
		mr_reg_attr.mr_iov = &iovec;
		mr_reg_attr.iov_count = 1;

		err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
		assert_int_equal(err, 0);

		/* FI_MR_DMABUF flag was not set, so GDRCopy should be registered if available */
		test_efa_rdm_mr_impl(efa_domain, mr, baseline_ct + 1, baseline_sz + mr_size, true);

		assert_int_equal(fi_close(&mr->fid), 0);
		test_efa_rdm_mr_impl(efa_domain, NULL, baseline_ct, baseline_sz, false);

		err = ofi_cudaFree(buf);
		assert_int_equal(err, 0);
	}
}
#else
void test_efa_rdm_mr_reg_cuda_memory(void **state)
{
	skip();
}
#endif

#if HAVE_CUDA
void test_efa_direct_mr_reg_cuda_memory(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	struct fi_mr_attr mr_reg_attr = { 0 };
	struct iovec iovec;
	int err, baseline_ct, baseline_sz;

	if (g_efa_hmem_info[FI_HMEM_CUDA].initialized &&
	    g_efa_hmem_info[FI_HMEM_CUDA].p2p_supported_by_device) {
		resource->hints = efa_unit_test_alloc_hints_hmem(
			FI_EP_RDM, EFA_DIRECT_FABRIC_NAME);
		efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
							    FI_VERSION(2, 0),
							    resource->hints,
							    true, true);

		efa_domain = container_of(resource->domain, struct efa_domain,
					  util_domain.domain_fid);
		/* fi_endpoint calls ofi_bufpool_grow, which registers mr */
		baseline_ct = ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct);
		baseline_sz = ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz);

		err = ofi_cudaMalloc(&buf, mr_size);
		assert_int_equal(err, 0);
		assert_non_null(buf);

		mr_reg_attr.access = FI_SEND | FI_RECV;
		mr_reg_attr.iface = FI_HMEM_CUDA;
		iovec.iov_base = buf;
		iovec.iov_len = mr_size;
		mr_reg_attr.mr_iov = &iovec;
		mr_reg_attr.iov_count = 1;

		err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
		assert_int_equal(err, 0);

		test_efa_mr_impl(efa_domain, mr, baseline_ct + 1, baseline_sz + mr_size);

		assert_int_equal(fi_close(&mr->fid), 0);
		test_efa_mr_impl(efa_domain, NULL, baseline_ct, baseline_sz);

		err = ofi_cudaFree(buf);
		assert_int_equal(err, 0);
	}
}
#else
void test_efa_direct_mr_reg_cuda_memory(void **state)
{
	skip();
}
#endif

void test_efa_direct_mr_reg_rdma_read_not_supported(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	struct fi_mr_attr mr_reg_attr = {0};
	struct iovec iovec;
	int err;
	uint32_t efa_device_caps_orig;

	resource->hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_DIRECT_FABRIC_NAME);
	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0),
						    resource->hints,
						    true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	/* Mock device to not support RDMA read */
	efa_device_caps_orig = efa_domain->device->device_caps;
	efa_domain->device->device_caps &= ~EFADV_DEVICE_ATTR_CAPS_RDMA_READ;

	buf = malloc(mr_size);
	assert_non_null(buf);

	mr_reg_attr.access = FI_READ;
	mr_reg_attr.iface = FI_HMEM_SYSTEM;
	iovec.iov_base = buf;
	iovec.iov_len = mr_size;
	mr_reg_attr.mr_iov = &iovec;
	mr_reg_attr.iov_count = 1;

	err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
	assert_int_equal(err, -FI_EOPNOTSUPP);

	free(buf);
	efa_domain->device->device_caps = efa_device_caps_orig;
}

void test_efa_direct_mr_reg_rdma_write_not_supported(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	struct fi_mr_attr mr_reg_attr = {0};
	struct iovec iovec;
	int err;
	uint32_t efa_device_caps_orig;

	resource->hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_DIRECT_FABRIC_NAME);
	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0),
						    resource->hints,
						    true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	efa_device_caps_orig = efa_domain->device->device_caps;
	/* Mock device to not support RDMA write */
#if HAVE_CAPS_RDMA_WRITE
	efa_domain->device->device_caps &= ~EFADV_DEVICE_ATTR_CAPS_RDMA_WRITE;
#endif

	buf = malloc(mr_size);
	assert_non_null(buf);

	mr_reg_attr.access = FI_WRITE;
	mr_reg_attr.iface = FI_HMEM_SYSTEM;
	iovec.iov_base = buf;
	iovec.iov_len = mr_size;
	mr_reg_attr.mr_iov = &iovec;
	mr_reg_attr.iov_count = 1;

	err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
	assert_int_equal(err, -FI_EOPNOTSUPP);

	free(buf);
	efa_domain->device->device_caps = efa_device_caps_orig;
}

void test_efa_mr_validate_regattr_invalid_iov_count(void **state)
{
	struct efa_resource *resource = *state;
	struct fi_mr_attr mr_attr = { 0 };
	struct iovec iov[2];
	void *buf1, *buf2;
	size_t mr_size = 64;
	struct fid_mr *mr = NULL;
	int ret;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);

	buf1 = malloc(mr_size);
	buf2 = malloc(mr_size);
	assert_non_null(buf1);
	assert_non_null(buf2);

	iov[0].iov_base = buf1;
	iov[0].iov_len = mr_size;
	iov[1].iov_base = buf2;
	iov[1].iov_len = mr_size;
	mr_attr.mr_iov = iov;
	mr_attr.iov_count = 2; /* Exceeds EFA_MR_IOV_LIMIT */
	mr_attr.access = FI_SEND | FI_RECV;
	mr_attr.iface = FI_HMEM_SYSTEM;

	/* Test through the actual API */
	ret = fi_mr_regattr(resource->domain, &mr_attr, 0, &mr);
	assert_int_equal(ret, -FI_EINVAL);
	assert_null(mr);

	free(buf1);
	free(buf2);
}

void test_efa_mr_validate_regattr_uninitialized_iface(void **state)
{
	struct efa_resource *resource = *state;
	struct fi_mr_attr mr_attr = { 0 };
	struct iovec iov;
	void *buf;
	size_t mr_size = 64;
	struct fid_mr *mr = NULL;
	int ret;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);

	buf = malloc(mr_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = mr_size;
	mr_attr.mr_iov = &iov;
	mr_attr.iov_count = 1;
	mr_attr.access = FI_SEND | FI_RECV;
	mr_attr.iface = FI_HMEM_CUDA;

	/* Mock CUDA as uninitialized by temporarily setting it to false */
	g_efa_hmem_info[FI_HMEM_CUDA].initialized = false;

	/* Test through the actual API */
	ret = fi_mr_regattr(resource->domain, &mr_attr, 0, &mr);
	assert_int_equal(ret, -FI_ENOSYS);
	assert_null(mr);

	free(buf);
}

/**
 * @brief Test RDM MR structure casting safety
 */
void test_efa_rdm_mr_structure_casting(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	struct efa_mr *efa_mr;
	struct efa_rdm_mr *efa_rdm_mr;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

				  buf = malloc(mr_size);
	assert_non_null(buf);

	assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
				   FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL),
			 0);

	/* Test safe casting from fid_mr to efa_mr */
	efa_mr = container_of(mr, struct efa_mr, mr_fid);
	assert_non_null(efa_mr);
	assert_ptr_equal(efa_mr->domain, efa_domain);
	assert_int_equal(efa_mr->iface, FI_HMEM_SYSTEM);

	/* Test safe casting from efa_mr to efa_rdm_mr */
	efa_rdm_mr = (struct efa_rdm_mr *)efa_mr;
	assert_non_null(efa_rdm_mr);
	/* Verify that efa_mr is the first member */
	assert_ptr_equal(&efa_rdm_mr->efa_mr, efa_mr);
	assert_ptr_equal(efa_rdm_mr->efa_mr.domain, efa_domain);
	assert_int_equal(efa_rdm_mr->efa_mr.iface, FI_HMEM_SYSTEM);

	assert_int_equal(fi_close(&mr->fid), 0);
	free(buf);
}

/**
 * @brief Test EFA_MR_ATTR_INIT_SYSTEM macro
 */
void test_efa_mr_attr_init_system_macro(void **state)
{
	struct iovec iov;
	void *buf;
	size_t mr_size = 64;
	uint64_t access = FI_SEND | FI_RECV;
	uint64_t offset = 0;
	uint64_t requested_key = 123;
	void *context = (void *)0xdeadbeef;

	buf = malloc(mr_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = mr_size;

	/* Test the macro by initializing a new struct */
	struct fi_mr_attr mr_attr = EFA_MR_ATTR_INIT_SYSTEM(&iov, 1, access, offset, requested_key, context);

	assert_ptr_equal(mr_attr.mr_iov, &iov);
	assert_int_equal(mr_attr.iov_count, 1);
	assert_int_equal(mr_attr.access, access);
	assert_int_equal(mr_attr.offset, offset);
	assert_int_equal(mr_attr.requested_key, requested_key);
	assert_ptr_equal(mr_attr.context, context);
	assert_int_equal(mr_attr.iface, FI_HMEM_SYSTEM);

	free(buf);
}

/**
 * @brief Test efa_mr_ofi_to_ibv_access with no access flags
 *
 * When no access flags are provided, the function should default to
 * FI_SEND | FI_RECV and return IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ.
 */
void test_efa_mr_ofi_to_ibv_access_no_access(void **state)
{
	int ibv_access;

	ibv_access = efa_mr_ofi_to_ibv_access(0, true, true);
	assert_int_equal(ibv_access, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ);
}

/**
 * @brief Test efa_mr_ofi_to_ibv_access with one flag when rdma read and write are available
 *
 */
void test_efa_mr_ofi_to_ibv_access_one_flag(void **state)
{
	int ibv_access;

	ibv_access = efa_mr_ofi_to_ibv_access(FI_SEND, true, true);
	assert_int_equal(ibv_access, IBV_ACCESS_REMOTE_READ);

	ibv_access = efa_mr_ofi_to_ibv_access(FI_RECV, true, true);
	assert_int_equal(ibv_access, IBV_ACCESS_LOCAL_WRITE);

	ibv_access = efa_mr_ofi_to_ibv_access(FI_READ, true, true);
	assert_int_equal(ibv_access, IBV_ACCESS_LOCAL_WRITE);

	ibv_access = efa_mr_ofi_to_ibv_access(FI_WRITE, true, true);
	assert_int_equal(ibv_access, 0);

	ibv_access = efa_mr_ofi_to_ibv_access(FI_REMOTE_READ, true, true);
	assert_int_equal(ibv_access, IBV_ACCESS_REMOTE_READ);

	ibv_access = efa_mr_ofi_to_ibv_access(FI_REMOTE_WRITE, true, true);
	assert_int_equal(ibv_access, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE);
}

/**
 * @brief Test efa_mr_ofi_to_ibv_access when RDMA read not supported
 */
void test_efa_mr_ofi_to_ibv_access_read_not_supported(void **state)
{
	int ibv_access;

	ibv_access = efa_mr_ofi_to_ibv_access(FI_READ, false, false);
	assert_int_equal(ibv_access, 0);

	ibv_access = efa_mr_ofi_to_ibv_access(FI_REMOTE_READ, false, false);
	assert_int_equal(ibv_access, 0);
}

/**
 * @brief Test efa_mr_ofi_to_ibv_access when RDMA write not supported
 *
 * When device doesn't support RDMA write, emulate with RDMA read
 */
void test_efa_mr_ofi_to_ibv_access_write_not_supported(void **state)
{
	int ibv_access;

	ibv_access = efa_mr_ofi_to_ibv_access(FI_WRITE, true, false);
	assert_int_equal(ibv_access, IBV_ACCESS_REMOTE_READ);

	ibv_access = efa_mr_ofi_to_ibv_access(FI_REMOTE_WRITE, true, false);
	assert_int_equal(ibv_access, IBV_ACCESS_LOCAL_WRITE);
}

/**
 * @brief Test efa_mr_ofi_to_ibv_access with FI_REMOTE_READ | FI_REMOTE_WRITE,
 * only read supported
 *
 * When only remote read is supported, FI_REMOTE_READ gets IBV_ACCESS_REMOTE_READ
 * and FI_REMOTE_WRITE gets IBV_ACCESS_LOCAL_WRITE.
 */
void test_efa_mr_ofi_to_ibv_access_remote_read_write_read_only_supported(void **state)
{
	int ibv_access;

	ibv_access = efa_mr_ofi_to_ibv_access(FI_REMOTE_READ | FI_REMOTE_WRITE, true, false);
	assert_int_equal(ibv_access, IBV_ACCESS_REMOTE_READ | IBV_ACCESS_LOCAL_WRITE);
}

/**
 * @brief Test efa_mr_ofi_to_ibv_access with all access flags combined
 *
 * Test all OFI access flags together with full device support.
 */
void test_efa_mr_ofi_to_ibv_access_all_flags_supported(void **state)
{
	int ibv_access;
	uint64_t all_flags = FI_SEND | FI_RECV | FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE;

	ibv_access = efa_mr_ofi_to_ibv_access(all_flags, true, true);
	assert_int_equal(ibv_access, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE);
}

/**
 * @brief Test efa_mr_ofi_to_ibv_access with all access flags, no device support
 */
void test_efa_mr_ofi_to_ibv_access_all_flags_not_supported(void **state)
{
	int ibv_access;
	uint64_t all_flags = FI_SEND | FI_RECV | FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE;

	ibv_access = efa_mr_ofi_to_ibv_access(all_flags, false, false);
	assert_int_equal(ibv_access, IBV_ACCESS_LOCAL_WRITE);
}

/**
 * @brief Test that closing an MR with outstanding direct operations prints
 * warning and clears the desc reference.
 *
 * When FI_EFA_TRACK_MR is enabled and an MR is closed while a direct
 * operation still references it, efa_mr_close will warn and clear the desc entry.
 */
void test_efa_mr_close_warn_outstanding_direct_ope(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_base_ep *base_ep;
	struct efa_direct_ope *direct_ope;
	struct fid_mr *mr = NULL;
	struct efa_mr *efa_mr;
	size_t mr_size = 64;
	void *buf;

	efa_env.track_mr = 1;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_DIRECT_FABRIC_NAME);

	base_ep = container_of(resource->ep, struct efa_base_ep,
			       util_ep.ep_fid);

	buf = malloc(mr_size);
	assert_non_null(buf);

	assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
				   FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL), 0);
	efa_mr = container_of(mr, struct efa_mr, mr_fid);

	/* Simulate an outstanding direct operation referencing this MR */
	direct_ope = ofi_buf_alloc(base_ep->efa_direct_ope_pool);
	assert_non_null(direct_ope);
	direct_ope->iov_count = 1;
	direct_ope->desc[0] = efa_mr;
	dlist_insert_tail(&direct_ope->entry, &base_ep->efa_direct_ope_list);

	/* Close MR while operation is outstanding */
	assert_int_equal(fi_close(&mr->fid), 0);

	/* Verify desc was cleared */
	assert_null(direct_ope->desc[0]);

	/* Clean up the outstanding ope */
	dlist_remove(&direct_ope->entry);
	ofi_buf_free(direct_ope);

	free(buf);

	/* Close ep before restoring track_mr so efa_ep_close sees track_mr=1
	 * and destroys the direct ope pool. */
	assert_int_equal(fi_close(&resource->ep->fid), 0);
	resource->ep = NULL;
}

/**
 * @brief Test MR close with outstanding direct operations on multiple EPs
 *
 * Two EPs share one MR. Each EP has an in-flight direct operation referencing
 * the MR. Closing the MR should warn and clear desc on both EPs.
 */
void test_efa_mr_close_warn_outstanding_direct_ope_multi_ep(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_base_ep *base_ep1, *base_ep2;
	struct efa_direct_ope *direct_ope1, *direct_ope2;
	struct fid_ep *ep2;
	struct fid_mr *mr = NULL;
	struct efa_mr *efa_mr;
	size_t mr_size = 64;
	void *buf;

	efa_env.track_mr = 1;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_DIRECT_FABRIC_NAME);

	base_ep1 = container_of(resource->ep, struct efa_base_ep,
				util_ep.ep_fid);

	/* Create and enable a second EP on the same domain and CQ */
	assert_int_equal(fi_endpoint(resource->domain, resource->info, &ep2, NULL), 0);
	assert_int_equal(fi_ep_bind(ep2, &resource->cq->fid, FI_SEND | FI_RECV), 0);
	assert_int_equal(fi_enable(ep2), 0);
	base_ep2 = container_of(ep2, struct efa_base_ep, util_ep.ep_fid);

	buf = malloc(mr_size);
	assert_non_null(buf);

	assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
				   FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL), 0);
	efa_mr = container_of(mr, struct efa_mr, mr_fid);

	/* Simulate outstanding direct operations on both EPs referencing the same MR */
	direct_ope1 = ofi_buf_alloc(base_ep1->efa_direct_ope_pool);
	assert_non_null(direct_ope1);
	direct_ope1->iov_count = 1;
	direct_ope1->desc[0] = efa_mr;
	dlist_insert_tail(&direct_ope1->entry, &base_ep1->efa_direct_ope_list);

	direct_ope2 = ofi_buf_alloc(base_ep2->efa_direct_ope_pool);
	assert_non_null(direct_ope2);
	direct_ope2->iov_count = 1;
	direct_ope2->desc[0] = efa_mr;
	dlist_insert_tail(&direct_ope2->entry, &base_ep2->efa_direct_ope_list);

	/* Close MR while operations are outstanding on both EPs */
	assert_int_equal(fi_close(&mr->fid), 0);

	/* Verify desc was cleared on both EPs */
	assert_null(direct_ope1->desc[0]);
	assert_null(direct_ope2->desc[0]);

	/* Clean up */
	dlist_remove(&direct_ope1->entry);
	ofi_buf_free(direct_ope1);
	dlist_remove(&direct_ope2->entry);
	ofi_buf_free(direct_ope2);

	assert_int_equal(fi_close(&ep2->fid), 0);
	free(buf);

	/* Close ep before restoring track_mr so efa_ep_close sees track_mr=1
	 * and destroys the direct ope pool. */
	assert_int_equal(fi_close(&resource->ep->fid), 0);
	resource->ep = NULL;
}

/**
 * @brief Test that closing an MR with outstanding RDM txe prints
 * warning and clears the desc reference.
 *
 * When FI_EFA_TRACK_MR is enabled and an MR is closed while an RDM
 * TX operation still references it, efa_mr_close will warn and clear the desc entry.
 */
void test_efa_mr_close_warn_outstanding_rdm_txe(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_rdm_ope *txe;
	struct fid_mr *mr = NULL;
	struct efa_mr *efa_mr;
	size_t mr_size = 64;
	void *buf;

	efa_env.track_mr = 1;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);

	buf = malloc(mr_size);
	assert_non_null(buf);

	assert_int_equal(fi_mr_reg(resource->domain, buf, mr_size,
				   FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL), 0);
	efa_mr = container_of(mr, struct efa_mr, mr_fid);

	/* Allocate a txe and set its desc to reference the MR */
	txe = efa_unit_test_alloc_txe(resource, ofi_op_msg);
	assert_non_null(txe);
	txe->iov_count = 1;
	txe->desc[0] = efa_mr;

	/* Close MR while txe is outstanding */
	assert_int_equal(fi_close(&mr->fid), 0);

	/* Verify desc was cleared */
	assert_null(txe->desc[0]);

	/* Clean up the txe */
	efa_rdm_txe_release(txe);

	free(buf);
}

/**
 * @brief Test efa_rdm_mr_cache_regv with cache disabled
 *
 * This test verifies that efa_rdm_mr_cache_regv properly handles the no-cache
 * scenario by falling back to internal registration and creates the expected
 * MR structure without cache entries.
 */
void test_efa_rdm_mr_cache_regv_no_cache(void **state)
{
	struct efa_resource *resource = *state;
	struct fi_info *hints;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	struct efa_rdm_mr *efa_rdm_mr;
	struct iovec iov;

	/* Create domain with cache disabled (FI_MR_LOCAL) */
	hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode |= FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0), hints, true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	/* Verify cache is not available */
	assert_null(efa_domain->cache);
	assert_false(efa_is_cache_available(efa_domain));

	buf = malloc(mr_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = mr_size;

	assert_int_equal(efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
					      FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL),
			 0);
	assert_non_null(mr);

	efa_rdm_mr = container_of(mr, struct efa_rdm_mr, efa_mr.mr_fid);
	/* Verify that shm_mr is NULL even if shm_domain exists */
	assert_null(efa_rdm_mr->shm_mr);

	/* Verify no cache entry since cache is disabled */
	assert_null(efa_rdm_mr->entry);
	/* Verify MR is inserted to domain map */
	assert_true(efa_rdm_mr->inserted_to_mr_map);

	assert_int_equal(fi_close(&mr->fid), 0);
	free(buf);
	fi_freeinfo(hints);
}

/**
 * @brief Test efa_rdm_mr_cache_regv with cache enabled
 *
 * This test validates that efa_rdm_mr_cache_regv properly uses the MR cache
 * when it's available and creates the expected MR structure.
 */
void test_efa_rdm_mr_cache_regv_with_cache(void **state)
{
#ifdef ENABLE_ASAN
	skip();
#else
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	struct fid_mr *mr = NULL;
	struct efa_rdm_mr *efa_rdm_mr;
	struct iovec iov;
	size_t mr_size = 64;
	void *buf;
	int ret;

	/* Create domain with cache enabled (no FI_MR_LOCAL) */
	struct fi_info *hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode &= ~FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0), hints, true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	/* Verify cache is available */
	assert_non_null(efa_domain->cache);
	assert_true(efa_is_cache_available(efa_domain));

	buf = malloc(mr_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = mr_size;

	/* Test efa_rdm_mr_cache_regv with cache available */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr);

	/* Verify it's an efa_rdm_mr from cache */
	efa_rdm_mr = container_of(mr, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr->entry); /* Should have cache entry */
	assert_true(efa_rdm_mr->inserted_to_mr_map);
	assert_ptr_equal(efa_rdm_mr->efa_mr.domain, efa_domain);

	assert_int_equal(fi_close(&mr->fid), 0);
	free(buf);
#endif
}

/**
 * @brief Test efa_rdm_mr_cache_regv cache hit scenario
 *
 * This test validates that multiple registrations of the same buffer
 * result in cache hits when cache is enabled.
 */
void test_efa_rdm_mr_cache_regv_cache_hit(void **state)
{
#ifdef ENABLE_ASAN
	skip();
#else
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	struct fid_mr *mr1 = NULL, *mr2 = NULL;
	struct efa_rdm_mr *efa_rdm_mr1, *efa_rdm_mr2;
	struct iovec iov;
	size_t mr_size = 64;
	void *buf;
	int ret;

	/* Create domain with cache enabled */
	struct fi_info *hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode &= ~FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0), hints, true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	/* Verify cache is available */
	assert_non_null(efa_domain->cache);

	buf = malloc(mr_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = mr_size;

	/* First registration - should create cache entry */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr1, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr1);

	efa_rdm_mr1 = container_of(mr1, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr1->entry);

	/* Second registration of same buffer - should hit cache */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr2, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr2);

	efa_rdm_mr2 = container_of(mr2, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr2->entry);

	/* Both should reference the same cache entry */
	assert_ptr_equal(efa_rdm_mr1->entry, efa_rdm_mr2->entry);

	assert_int_equal(fi_close(&mr1->fid), 0);
	assert_int_equal(fi_close(&mr2->fid), 0);
	free(buf);
#endif
}

/**
 * @brief Test MR cache encapsulation behavior - smaller region within larger cached region
 *
 * New registration is smaller and fully encapsulated within a previous region.
 * The reference count should be incremented and the same cache entry returned.
 */
void test_efa_rdm_mr_cache_encapsulation_smaller(void **state)
{
#ifdef ENABLE_ASAN
	skip();
#else
	struct efa_resource *resource = *state;
	struct fid_mr *mr_large = NULL, *mr_small = NULL;
	struct efa_rdm_mr *efa_rdm_mr_large, *efa_rdm_mr_small;
	struct iovec iov_large, iov_small;
	size_t large_size = 1024, small_size = 256;
	void *buf;
	int ret;

	/* Create domain with cache enabled */
	struct fi_info *hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode &= ~FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0), hints, true, true);

	buf = malloc(large_size);
	assert_non_null(buf);

	/* Register large region first */
	iov_large.iov_base = buf;
	iov_large.iov_len = large_size;

	ret = efa_rdm_mr_cache_regv(resource->domain, &iov_large, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr_large, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr_large);

	efa_rdm_mr_large = container_of(mr_large, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr_large->entry);

	/* Register smaller region within the large region */
	iov_small.iov_base = (char *)buf + 100; /* Offset within large region */
	iov_small.iov_len = small_size;

	ret = efa_rdm_mr_cache_regv(resource->domain, &iov_small, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr_small, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr_small);

	efa_rdm_mr_small = container_of(mr_small, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr_small->entry);

	/* Both should reference the same cache entry (encapsulation) */
	assert_ptr_equal(efa_rdm_mr_large->entry, efa_rdm_mr_small->entry);

	assert_int_equal(fi_close(&mr_large->fid), 0);
	assert_int_equal(fi_close(&mr_small->fid), 0);
	free(buf);
#endif
}

/**
 * @brief Test MR cache non-overlapping regions behavior
 *
 * New registration and previous registrations have no full encapsulations.
 * A new registration should be created with reference count of 1.
 */
void test_efa_rdm_mr_cache_non_overlapping(void **state)
{
#ifdef ENABLE_ASAN
	skip();
#else
	struct efa_resource *resource = *state;
	struct fid_mr *mr1 = NULL, *mr2 = NULL;
	struct efa_rdm_mr *efa_rdm_mr1, *efa_rdm_mr2;
	struct iovec iov1, iov2;
	size_t mr_size = 256;
	void *buf1, *buf2;
	int ret;

	/* Create domain with cache enabled */
	struct fi_info *hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode &= ~FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0), hints, true, true);

	/* Allocate two separate, non-overlapping buffers */
	buf1 = malloc(mr_size);
	buf2 = malloc(mr_size);
	assert_non_null(buf1);
	assert_non_null(buf2);

	/* Ensure buffers are non-overlapping by checking addresses */
	assert_true((char *)buf1 + mr_size <= (char *)buf2 ||
		    (char *)buf2 + mr_size <= (char *)buf1);

	/* Register first region */
	iov1.iov_base = buf1;
	iov1.iov_len = mr_size;

	ret = efa_rdm_mr_cache_regv(resource->domain, &iov1, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr1, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr1);

	efa_rdm_mr1 = container_of(mr1, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr1->entry);

	/* Register second, non-overlapping region */
	iov2.iov_base = buf2;
	iov2.iov_len = mr_size;

	ret = efa_rdm_mr_cache_regv(resource->domain, &iov2, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr2, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr2);

	efa_rdm_mr2 = container_of(mr2, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr2->entry);

	/* Should have different cache entries (non-overlapping) */
	assert_ptr_not_equal(efa_rdm_mr1->entry, efa_rdm_mr2->entry);

	assert_int_equal(fi_close(&mr1->fid), 0);
	assert_int_equal(fi_close(&mr2->fid), 0);
	free(buf1);
	free(buf2);
#endif
}

/**
 * @brief Test MR cache LRU behavior
 *
 * Tests that MRs with reference count of zero are moved to LRU list
 * and can be reused when the same region is registered again.
 */
void test_efa_rdm_mr_cache_lru_behavior(void **state)
{
#ifdef ENABLE_ASAN
	skip();
#else
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	struct fid_mr *mr1 = NULL, *mr2 = NULL;
	struct efa_rdm_mr *efa_rdm_mr1, *efa_rdm_mr2;
	struct iovec iov;
	size_t mr_size = 256;
	void *buf;
	int ret;

	/* Create domain with cache enabled */
	struct fi_info *hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode &= ~FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0), hints, true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	buf = malloc(mr_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = mr_size;

	/* Initially LRU list should be empty */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* First registration */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr1, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr1);

	efa_rdm_mr1 = container_of(mr1, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr1->entry);
	/* LRU should be empty while entry is in use */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* Close first MR - should move to LRU list */
	assert_int_equal(fi_close(&mr1->fid), 0);
	/* LRU should now contain the entry */
	assert_false(dlist_empty(&efa_domain->cache->lru_list));

	/* Register same region again - should reuse from LRU */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr2, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr2);

	efa_rdm_mr2 = container_of(mr2, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr2->entry);

	/* Should reuse the same cache entry from LRU */
	assert_ptr_equal(efa_rdm_mr1->entry, efa_rdm_mr2->entry);
	/* LRU should be empty again since entry is back in use */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	assert_int_equal(fi_close(&mr2->fid), 0);
	/* LRU should contain the entry again */
	assert_false(dlist_empty(&efa_domain->cache->lru_list));

	free(buf);
#endif
}

/**
 * @brief Test MR cache flush behavior
 *
 * Tests that cache flush operations properly clean up entries
 * and that subsequent registrations create new entries.
 */
void test_efa_rdm_mr_cache_flush_behavior(void **state)
{
#ifdef ENABLE_ASAN
	skip();
#else
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	struct fid_mr *mr1 = NULL, *mr2 = NULL;
	struct efa_rdm_mr *efa_rdm_mr1, *efa_rdm_mr2;
	struct iovec iov;
	size_t mr_size = 256;
	void *buf;
	int ret;

	/* Create domain with cache enabled */
	struct fi_info *hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode &= ~FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0), hints, true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	buf = malloc(mr_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = mr_size;

	/* Initially LRU list should be empty */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* First registration */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr1, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr1);

	efa_rdm_mr1 = container_of(mr1, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr1->entry);
	/* LRU should be empty while entry is in use */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* Close MR - moves to LRU */
	assert_int_equal(fi_close(&mr1->fid), 0);
	/* LRU should now contain the entry */
	assert_false(dlist_empty(&efa_domain->cache->lru_list));

	/* Force cache flush - should clean up LRU entries */
	ofi_mr_cache_flush(efa_domain->cache, true /* flush_lru */);
	/* LRU should be empty after flush */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* Register same region again - should create new entry after flush */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr2, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr2);

	efa_rdm_mr2 = container_of(mr2, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr2->entry);
	/* Should have a valid new cache entry */
	assert_int_equal(efa_rdm_mr2->entry->use_cnt, 1);

	assert_int_equal(fi_close(&mr2->fid), 0);
	free(buf);
#endif
}

/**
 * @brief Test MR cache reference counting
 *
 * Tests that multiple references to the same cached region properly
 * increment/decrement reference counts and manage LRU list placement.
 */
void test_efa_rdm_mr_cache_reference_counting(void **state)
{
#ifdef ENABLE_ASAN
	skip();
#else
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	struct fid_mr *mr1 = NULL, *mr2 = NULL, *mr3 = NULL;
	struct efa_rdm_mr *efa_rdm_mr1, *efa_rdm_mr2, *efa_rdm_mr3;
	struct iovec iov;
	size_t mr_size = 256;
	void *buf;
	int ret;

	/* Create domain with cache enabled */
	struct fi_info *hints = efa_unit_test_alloc_hints(FI_EP_RDM, EFA_FABRIC_NAME);
	hints->domain_attr->mr_mode &= ~FI_MR_LOCAL;

	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0), hints, true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	buf = malloc(mr_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = mr_size;

	/* Initially LRU list should be empty */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* First registration - creates cache entry with use_cnt = 1 */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr1, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr1);

	efa_rdm_mr1 = container_of(mr1, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr1->entry);
	assert_int_equal(efa_rdm_mr1->entry->use_cnt, 1);
	/* LRU should still be empty since use_cnt > 0 */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* Second registration - should increment reference count to 2 */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr2, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr2);

	efa_rdm_mr2 = container_of(mr2, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_ptr_equal(efa_rdm_mr1->entry, efa_rdm_mr2->entry);
	assert_int_equal(efa_rdm_mr1->entry->use_cnt, 2);
	/* LRU should still be empty since use_cnt > 0 */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* Third registration - should increment reference count to 3 */
	ret = efa_rdm_mr_cache_regv(resource->domain, &iov, 1,
				    FI_SEND | FI_RECV, 0, 0, 0, &mr3, NULL);
	assert_int_equal(ret, 0);
	assert_non_null(mr3);

	efa_rdm_mr3 = container_of(mr3, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_ptr_equal(efa_rdm_mr1->entry, efa_rdm_mr3->entry);
	assert_int_equal(efa_rdm_mr1->entry->use_cnt, 3);
	/* LRU should still be empty since use_cnt > 0 */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* Close first MR - should decrement reference count to 2, not move to LRU */
	assert_int_equal(fi_close(&mr1->fid), 0);
	assert_int_equal(efa_rdm_mr2->entry->use_cnt, 2);
	/* LRU should still be empty since use_cnt > 0 */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* Close second MR - should decrement reference count to 1, not move to LRU */
	assert_int_equal(fi_close(&mr2->fid), 0);
	assert_int_equal(efa_rdm_mr3->entry->use_cnt, 1);
	/* LRU should still be empty since use_cnt > 0 */
	assert_true(dlist_empty(&efa_domain->cache->lru_list));

	/* Close third MR - should decrement to zero and move to LRU */
	assert_int_equal(fi_close(&mr3->fid), 0);
	/* Now LRU should contain the entry since use_cnt reached 0 */
	assert_false(dlist_empty(&efa_domain->cache->lru_list));

	free(buf);
#endif
}

#if HAVE_CUDA
/**
 * @brief Test RDM MR registration for CUDA memory in non-p2p path
 *
 * This test verifies that when p2p is not supported, the non-p2p path
 * properly initializes the MR structure, including the mem_desc field
 * and iface field. It uses fi_mr_desc to verify the descriptor is correct.
 */
void test_efa_rdm_mr_reg_cuda_memory_non_p2p(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_domain *efa_domain;
	size_t mr_size = 64;
	void *buf;
	struct fid_mr *mr = NULL;
	struct fi_mr_attr mr_reg_attr = { 0 };
	struct iovec iovec;
	struct efa_rdm_mr *efa_rdm_mr;
	void *desc;
	int err, baseline_ct, baseline_sz;

	if (!g_efa_hmem_info[FI_HMEM_CUDA].initialized) {
		skip();
		return;
	}

	resource->hints = efa_unit_test_alloc_hints_hmem(
		FI_EP_RDM, EFA_FABRIC_NAME);
	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0),
						    resource->hints,
						    true, true);

	efa_domain = container_of(resource->domain, struct efa_domain,
				  util_domain.domain_fid);

	/* Mock p2p as not supported to force non-p2p path */
	g_efa_hmem_info[FI_HMEM_CUDA].p2p_supported_by_device = false;

	/* fi_endpoint calls ofi_bufpool_grow, which registers mr */
	baseline_ct = ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct);
	baseline_sz = ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz);

	err = ofi_cudaMalloc(&buf, mr_size);
	assert_int_equal(err, 0);
	assert_non_null(buf);

	mr_reg_attr.access = FI_SEND | FI_RECV;
	mr_reg_attr.iface = FI_HMEM_CUDA;
	iovec.iov_base = buf;
	iovec.iov_len = mr_size;
	mr_reg_attr.mr_iov = &iovec;
	mr_reg_attr.iov_count = 1;

	err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
	assert_int_equal(err, 0);
	assert_non_null(mr);

	/* Verify this is an RDM MR */
	efa_rdm_mr = container_of(mr, struct efa_rdm_mr, efa_mr.mr_fid);
	assert_non_null(efa_rdm_mr);

	/* Test fi_mr_desc returns the correct descriptor */
	desc = fi_mr_desc(mr);
	assert_non_null(desc);
	assert_ptr_equal(desc, &efa_rdm_mr->efa_mr);

	/* Verify iface is properly set (second bug fix) */
	assert_int_equal(efa_rdm_mr->efa_mr.iface, FI_HMEM_CUDA);

	/* Verify key is set (non-p2p path should generate proprietary key) */
	assert_true(fi_mr_key(mr) != FI_KEY_NOTAVAIL);
	assert_true(fi_mr_key(mr) > UINT32_MAX); /* Non-p2p keys should be > UINT32_MAX */

	/* Since we're in non-p2p path, no ibv_mr should be registered */
	assert_int_equal(ofi_atomic_get64(&efa_domain->ibv_mr_reg_ct), baseline_ct);
	assert_int_equal(ofi_atomic_get64(&efa_domain->ibv_mr_reg_sz), baseline_sz);

	/* Verify RDM-specific fields are properly initialized */
	assert_true(efa_rdm_mr->inserted_to_mr_map);
	assert_true(efa_rdm_mr->needs_sync); /* CUDA memory should need sync */

	/* Cleanup */
	assert_int_equal(fi_close(&mr->fid), 0);
	err = ofi_cudaFree(buf);
	assert_int_equal(err, 0);
}
#else
void test_efa_rdm_mr_reg_cuda_memory_non_p2p(void **state)
{
	skip();
}
#endif

/**
 * @brief Verify that fi_mr_regattr rejects an out-of-range iface value
 *
 * @param[in]	state		struct efa_resource that is managed by the framework
 */
void test_efa_mr_reg_out_of_range_iface(void **state)
{
	struct efa_resource *resource = *state;
	struct fi_mr_attr mr_reg_attr = {0};
	struct iovec iov;
	struct fid_mr *mr = NULL;
	char buf[64];
	int err;

	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);

	iov.iov_base = buf;
	iov.iov_len = sizeof(buf);
	mr_reg_attr.mr_iov = &iov;
	mr_reg_attr.iov_count = 1;
	mr_reg_attr.access = FI_SEND | FI_RECV;

	/* Test with value >= OFI_HMEM_MAX */
	mr_reg_attr.iface = OFI_HMEM_MAX + 100;
	err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
	assert_int_equal(err, -FI_EINVAL);

	/* Test with negative value */
	mr_reg_attr.iface = (enum fi_hmem_iface)(-1);
	err = fi_mr_regattr(resource->domain, &mr_reg_attr, 0, &mr);
	assert_int_equal(err, -FI_EINVAL);
	assert_null(mr);
}

/**
 * Verify that efa_direct_ope is released when fi_recv fails after
 * allocating the ope. Without the fix the ope leaks on the error path.
 */
void test_efa_direct_ope_released_on_recv_error(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_base_ep *base_ep;
	struct efa_context ctx = {0};
	size_t buf_size = 64;
	void *buf;
	struct iovec iov;
	struct fi_msg msg = {0};
	int ret;

	efa_env.track_mr = 1;

	resource->hints = efa_unit_test_alloc_hints(FI_EP_RDM,
						    EFA_DIRECT_FABRIC_NAME);
	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0),
						    resource->hints,
						    true, true);

	base_ep = container_of(resource->ep, struct efa_base_ep,
			       util_ep.ep_fid);

	buf = malloc(buf_size);
	assert_non_null(buf);

	/* Post a recv with NULL desc to trigger -FI_EINVAL after ope alloc */
	iov.iov_base = buf;
	iov.iov_len = buf_size;
	msg.msg_iov = &iov;
	msg.iov_count = 1;
	msg.desc = NULL;
	msg.context = &ctx;

	ret = fi_recvmsg(resource->ep, &msg, FI_COMPLETION);
	assert_int_equal(ret, -FI_EINVAL);

	/* The ope list must be empty - ope released on error path */
	assert_true(dlist_empty(&base_ep->efa_direct_ope_list));

	free(buf);
}

/**
 * Verify that efa_direct_ope is released when fi_sendmsg fails after
 * allocating the ope. Without the fix the ope leaks on the error path.
 */
void test_efa_direct_ope_released_on_send_error(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_base_ep *base_ep;
	struct efa_context ctx = {0};
	struct efa_ep_addr raw_addr;
	size_t raw_addr_len = sizeof(raw_addr);
	fi_addr_t peer_addr;
	size_t buf_size = 64;
	void *buf;
	struct iovec iov;
	struct fi_msg msg = {0};
	int ret;

	efa_env.track_mr = 1;

	resource->hints = efa_unit_test_alloc_hints(FI_EP_RDM,
						    EFA_DIRECT_FABRIC_NAME);
	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0),
						    resource->hints,
						    true, true);

	base_ep = container_of(resource->ep, struct efa_base_ep,
			       util_ep.ep_fid);

	/* Insert a fake peer so efa_av_addr_to_conn succeeds */
	assert_int_equal(fi_getname(&resource->ep->fid, &raw_addr,
				    &raw_addr_len), 0);
	raw_addr.qpn = 0;
	raw_addr.qkey = 0x1234;
	assert_int_equal(fi_av_insert(resource->av, &raw_addr, 1,
				      &peer_addr, 0, NULL), 1);

	buf = malloc(buf_size);
	assert_non_null(buf);

	/* Post a send with NULL desc to trigger -FI_EINVAL after ope alloc */
	iov.iov_base = buf;
	iov.iov_len = buf_size;
	msg.msg_iov = &iov;
	msg.iov_count = 1;
	msg.desc = NULL;
	msg.context = &ctx;
	msg.addr = peer_addr;

	ret = fi_sendmsg(resource->ep, &msg, FI_COMPLETION);
	assert_int_equal(ret, -FI_EINVAL);

	/* The ope list must be empty - ope released on error path */
	assert_true(dlist_empty(&base_ep->efa_direct_ope_list));

	free(buf);
}

/**
 * Verify that efa_direct_ope is released when fi_readmsg fails after
 * allocating the ope. Without the fix the ope leaks on the error path.
 */
void test_efa_direct_ope_released_on_read_error(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_base_ep *base_ep;
	struct efa_context ctx = {0};
	struct efa_ep_addr raw_addr;
	size_t raw_addr_len = sizeof(raw_addr);
	fi_addr_t peer_addr;
	size_t buf_size = 64;
	void *buf;
	struct iovec iov;
	struct fi_rma_iov rma_iov = {0};
	struct fi_msg_rma msg = {0};
	int ret;

	/* Skip test on platforms that don't support RDMA. */
	if (!efa_device_support_rdma_read() || !efa_device_support_rdma_write()) {
		skip();
		return;
	}

	efa_env.track_mr = 1;

	resource->hints = efa_unit_test_alloc_hints(FI_EP_RDM,
						    EFA_DIRECT_FABRIC_NAME);
	resource->hints->caps |= FI_RMA;
	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0),
						    resource->hints,
						    true, true);

	base_ep = container_of(resource->ep, struct efa_base_ep,
			       util_ep.ep_fid);

	assert_int_equal(fi_getname(&resource->ep->fid, &raw_addr,
				    &raw_addr_len), 0);
	raw_addr.qpn = 0;
	raw_addr.qkey = 0x1234;
	assert_int_equal(fi_av_insert(resource->av, &raw_addr, 1,
				      &peer_addr, 0, NULL), 1);

	buf = malloc(buf_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = buf_size;
	rma_iov.addr = 0x1000;
	rma_iov.len = buf_size;
	rma_iov.key = 0;
	msg.msg_iov = &iov;
	msg.iov_count = 1;
	msg.desc = NULL;
	msg.rma_iov = &rma_iov;
	msg.rma_iov_count = 1;
	msg.context = &ctx;
	msg.addr = peer_addr;

	ret = fi_readmsg(resource->ep, &msg, FI_COMPLETION);
	assert_int_equal(ret, -FI_EINVAL);

	assert_true(dlist_empty(&base_ep->efa_direct_ope_list));

	free(buf);
}

/**
 * Verify that efa_direct_ope is released when fi_writemsg fails after
 * allocating the ope. Without the fix the ope leaks on the error path.
 */
void test_efa_direct_ope_released_on_write_error(void **state)
{
	struct efa_resource *resource = *state;
	struct efa_base_ep *base_ep;
	struct efa_context ctx = {0};
	struct efa_ep_addr raw_addr;
	size_t raw_addr_len = sizeof(raw_addr);
	fi_addr_t peer_addr;
	size_t buf_size = 64;
	void *buf;
	struct iovec iov;
	struct fi_rma_iov rma_iov = {0};
	struct fi_msg_rma msg = {0};
	int ret;

	/* Skip test on platforms that don't support RDMA. */
	if (!efa_device_support_rdma_read() || !efa_device_support_rdma_write()) {
		skip();
		return;
	}

	efa_env.track_mr = 1;

	resource->hints = efa_unit_test_alloc_hints(FI_EP_RDM,
						    EFA_DIRECT_FABRIC_NAME);
	resource->hints->caps |= FI_RMA;
	efa_unit_test_resource_construct_with_hints(resource, FI_EP_RDM,
						    FI_VERSION(2, 0),
						    resource->hints,
						    true, true);

	base_ep = container_of(resource->ep, struct efa_base_ep,
			       util_ep.ep_fid);

	assert_int_equal(fi_getname(&resource->ep->fid, &raw_addr,
				    &raw_addr_len), 0);
	raw_addr.qpn = 0;
	raw_addr.qkey = 0x1234;
	assert_int_equal(fi_av_insert(resource->av, &raw_addr, 1,
				      &peer_addr, 0, NULL), 1);

	buf = malloc(buf_size);
	assert_non_null(buf);

	iov.iov_base = buf;
	iov.iov_len = buf_size;
	rma_iov.addr = 0x1000;
	rma_iov.len = buf_size;
	rma_iov.key = 0;
	msg.msg_iov = &iov;
	msg.iov_count = 1;
	msg.desc = NULL;
	msg.rma_iov = &rma_iov;
	msg.rma_iov_count = 1;
	msg.context = &ctx;
	msg.addr = peer_addr;

	ret = fi_writemsg(resource->ep, &msg, FI_COMPLETION);
	assert_int_equal(ret, -FI_EINVAL);

	assert_true(dlist_empty(&base_ep->efa_direct_ope_list));

	free(buf);
}
