diff --git a/src/pdr.c b/src/pdr.c
index ff4fbd1..752fd57 100644
--- a/src/pdr.c
+++ b/src/pdr.c
@@ -447,8 +447,22 @@
 	uint16_t entity_instance_number, pldm_entity_node *parent,
 	uint8_t association_type)
 {
-	assert(tree != NULL);
-	assert(entity != NULL);
+	return pldm_entity_association_tree_add_entity(tree, entity,
+						       entity_instance_number,
+						       parent, association_type,
+						       false, true, 0xFFFF);
+}
+
+LIBPLDM_ABI_TESTING
+pldm_entity_node *pldm_entity_association_tree_add_entity(
+	pldm_entity_association_tree *tree, pldm_entity *entity,
+	uint16_t entity_instance_number, pldm_entity_node *parent,
+	uint8_t association_type, bool is_remote, bool is_update_container_id,
+	uint16_t container_id)
+{
+	if ((!tree) || (!entity)) {
+		return NULL;
+	}
 
 	if (entity_instance_number != 0xFFFF && parent != NULL) {
 		pldm_entity node;
@@ -458,11 +472,14 @@
 			return NULL;
 		}
 	}
-
-	assert(association_type == PLDM_ENTITY_ASSOCIAION_PHYSICAL ||
-	       association_type == PLDM_ENTITY_ASSOCIAION_LOGICAL);
+	if (association_type != PLDM_ENTITY_ASSOCIAION_PHYSICAL &&
+	    association_type != PLDM_ENTITY_ASSOCIAION_LOGICAL) {
+		return NULL;
+	}
 	pldm_entity_node *node = malloc(sizeof(pldm_entity_node));
-	assert(node != NULL);
+	if (!node) {
+		return NULL;
+	}
 	node->first_child = NULL;
 	node->next_sibling = NULL;
 	node->parent.entity_type = 0;
@@ -473,9 +490,11 @@
 		entity_instance_number != 0xFFFF ? entity_instance_number : 1;
 	node->association_type = association_type;
 	node->remote_container_id = 0;
-
 	if (tree->root == NULL) {
-		assert(parent == NULL);
+		if (parent != NULL) {
+			free(node);
+			return NULL;
+		}
 		tree->root = node;
 		/* container_id 0 here indicates this is the top-most entry */
 		node->entity.entity_container_id = 0;
@@ -483,16 +502,41 @@
 	} else if (parent != NULL && parent->first_child == NULL) {
 		parent->first_child = node;
 		node->parent = parent->entity;
-		node->entity.entity_container_id = next_container_id(tree);
+
+		if (is_remote) {
+			node->remote_container_id = entity->entity_container_id;
+		}
+		if (is_update_container_id) {
+			if (container_id != 0xFFFF) {
+				node->entity.entity_container_id = container_id;
+			} else {
+				node->entity.entity_container_id =
+					next_container_id(tree);
+			}
+		} else {
+			node->entity.entity_container_id =
+				entity->entity_container_id;
+		}
+
+		if (!is_remote) {
+			node->remote_container_id =
+				node->entity.entity_container_id;
+		}
 	} else {
 		pldm_entity_node *start = parent == NULL ? tree->root :
 							   parent->first_child;
 		pldm_entity_node *prev =
 			find_insertion_at(start, entity->entity_type);
-		assert(prev != NULL);
+		if (!prev) {
+			free(node);
+			return NULL;
+		}
 		pldm_entity_node *next = prev->next_sibling;
 		if (prev->entity.entity_type == entity->entity_type) {
-			assert(prev->entity.entity_instance_num != UINT16_MAX);
+			if (prev->entity.entity_instance_num == UINT16_MAX) {
+				free(node);
+				return NULL;
+			}
 			node->entity.entity_instance_num =
 				entity_instance_number != 0xFFFF ?
 					entity_instance_number :
@@ -506,8 +550,9 @@
 		node->remote_container_id = entity->entity_container_id;
 	}
 	entity->entity_instance_num = node->entity.entity_instance_num;
-	entity->entity_container_id = node->entity.entity_container_id;
-
+	if (is_update_container_id) {
+		entity->entity_container_id = node->entity.entity_container_id;
+	}
 	return node;
 }
 
@@ -921,6 +966,51 @@
 	}
 }
 
+static void entity_association_tree_find_if_remote(pldm_entity_node *node,
+						   pldm_entity *entity,
+						   pldm_entity_node **out,
+						   bool is_remote)
+{
+	assert(out != NULL && *out == NULL);
+	if (node == NULL) {
+		return;
+	}
+	bool is_entity_type;
+	bool is_entity_instance_num;
+
+	is_entity_type = node->entity.entity_type == entity->entity_type;
+	is_entity_instance_num = node->entity.entity_instance_num ==
+				 entity->entity_instance_num;
+
+	if (!is_remote ||
+	    node->remote_container_id == entity->entity_container_id) {
+		if (is_entity_type && is_entity_instance_num) {
+			entity->entity_container_id =
+				node->entity.entity_container_id;
+			*out = node;
+			return;
+		}
+	}
+	entity_association_tree_find_if_remote(node->next_sibling, entity, out,
+					       is_remote);
+	entity_association_tree_find_if_remote(node->first_child, entity, out,
+					       is_remote);
+}
+
+LIBPLDM_ABI_TESTING
+pldm_entity_node *
+pldm_entity_association_tree_find_if_remote(pldm_entity_association_tree *tree,
+					    pldm_entity *entity, bool is_remote)
+{
+	if (!tree || !entity) {
+		return NULL;
+	}
+	pldm_entity_node *node = NULL;
+	entity_association_tree_find_if_remote(tree->root, entity, &node,
+					       is_remote);
+	return node;
+}
+
 LIBPLDM_ABI_STABLE
 void entity_association_tree_find(pldm_entity_node *node, pldm_entity *entity,
 				  pldm_entity_node **out)