Prefab | Improve create prefab code in public handler (#12506)

* Improve create prefab code in public handler

Signed-off-by: Junhao Wang <[email protected]>

* Use container entity DOM as a more selective approach

Signed-off-by: Junhao Wang <[email protected]>

* Add comment about the prefab dom value move

Signed-off-by: Junhao Wang <[email protected]>

* Removes the reset logic and make DOM updates selective

Signed-off-by: Junhao Wang <[email protected]>

* Update CreatePrefabInDisk function name

Signed-off-by: Junhao Wang <[email protected]>

* Revert API changes

Signed-off-by: Junhao Wang <[email protected]>

* Remove capturing intermediate dom state and add common root check

Signed-off-by: Junhao Wang <[email protected]>

Signed-off-by: Junhao Wang <[email protected]>

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabPublicHandler.cpp

@@ -95,6 +95,8 @@ namespace AzToolsFramework
             EntityList inputEntityList, topLevelEntities;
             AZ::EntityId commonRootEntityId;
             InstanceOptionalReference commonRootEntityOwningInstance;
+
+            // Find common root entity.
             PrefabOperationResult findCommonRootOutcome = FindCommonRootOwningInstance(
                 entityIds, inputEntityList, topLevelEntities, commonRootEntityId, commonRootEntityOwningInstance);
             if (!findCommonRootOutcome.IsSuccess())
@@ -102,7 +104,9 @@ namespace AzToolsFramework
                 return AZ::Failure(findCommonRootOutcome.TakeError());
             }
 
-            // order entities by their respective position within Entity Outliner
+            const TemplateId commonRootInstanceTemplateId = commonRootEntityOwningInstance->get().GetTemplateId();
+
+            // Order entities by their respective positions within hierarchy.
             EditorEntitySortRequestBus::Event(
                 commonRootEntityId,
                 [&topLevelEntities](EditorEntitySortRequestBus::Events* sortRequests)
@@ -116,31 +120,32 @@ namespace AzToolsFramework
                         });
                 });
 
-            AZ::EntityId containerEntityId;
-
+            AZ::EntityId newContainerEntityId;
             InstanceOptionalReference instanceToCreate;
             {
                 // Initialize Undo Batch object
                 ScopedUndoBatch undoBatch("Create Prefab");
 
-                PrefabDom commonRootInstanceDomBeforeCreate;
-                m_instanceToTemplateInterface->GenerateDomForInstance(
-                    commonRootInstanceDomBeforeCreate, commonRootEntityOwningInstance->get());
-
-                AZStd::vector<AZ::Entity*> entities;
-                AZStd::vector<AZStd::unique_ptr<Instance>> instancePtrs;
-                AZStd::vector<Instance*> instances;
+                AZStd::vector<AZ::Entity*> detachedEntities;
+                AZStd::vector<Instance*> detachedInstances;
+                AZStd::vector<AZStd::unique_ptr<Instance>> detachedInstancePtrs;
                 AZStd::unordered_map<Instance*, PrefabDom> nestedInstanceLinkPatchesMap;
 
                 // Retrieve all entities affected and identify Instances
                 PrefabOperationResult retrieveEntitiesAndInstancesOutcome = RetrieveAndSortPrefabEntitiesAndInstances(
-                    inputEntityList, commonRootEntityOwningInstance->get(), entities, instances);
+                    inputEntityList, commonRootEntityOwningInstance->get(), detachedEntities, detachedInstances);
                 if (!retrieveEntitiesAndInstancesOutcome.IsSuccess())
                 {
                     return AZ::Failure(retrieveEntitiesAndInstancesOutcome.TakeError());
                 }
 
-                AZStd::unordered_map<AZ::EntityId, AZStd::string> oldEntityAliases;
+                // Capture parent DOM state before detaching.
+                PrefabDom commonRootEntityDomBeforeDetaching;
+                if (commonRootEntityId.IsValid())
+                {
+                    m_instanceToTemplateInterface->GenerateDomForEntity(commonRootEntityDomBeforeDetaching,
+                        *GetEntityById(commonRootEntityId));
+                }
 
                 // Calculate new transform data before updating the direct top level entities.
                 AZ::Vector3 containerEntityTranslation(AZ::Vector3::CreateZero());
@@ -155,36 +160,54 @@ namespace AzToolsFramework
                 }
 
                 // Detach the retrieved entities. This includes the top level entities.
-                for (AZ::Entity* entity : entities)
+                AZStd::unordered_map<AZ::EntityId, AZStd::string> oldEntityAliases;
                 {
-                    AZ::EntityId entityId = entity->GetId();
-                    oldEntityAliases.emplace(entityId, commonRootEntityOwningInstance->get().GetEntityAlias(entityId)->get());
-                    commonRootEntityOwningInstance->get().DetachEntity(entity->GetId()).release();
-                }
+                    // DOM value pointers can't be relied upon if the original DOM gets modified after pointer creation.
+                    // This scope is added to limit their usage and ensure DOM is not modified when it is being used.
+                    AZStd::vector<AZStd::pair<const PrefabDomValue*, AZStd::string>> detachedEntityDomAndPathList;
+                    const PrefabDomValue& commonRootInstanceTemplateDom = m_prefabSystemComponentInterface->FindTemplateDom(
+                        commonRootInstanceTemplateId);
 
-                PrefabUndoHelpers::UpdatePrefabInstance(
-                    commonRootEntityOwningInstance->get(), "Update prefab instance", commonRootInstanceDomBeforeCreate,
-                    undoBatch.GetUndoBatch());
+                    for (AZ::Entity* entity : detachedEntities)
+                    {
+                        const AZ::EntityId entityId = entity->GetId();
+                        const AZStd::string entityAliasPath = m_instanceToTemplateInterface->GenerateEntityAliasPath(entityId);
+
+                        PrefabDomPath entityDomPathInOwningTemplate(entityAliasPath.c_str());
+                        const PrefabDomValue* detachedEntityDom = entityDomPathInOwningTemplate.Get(commonRootInstanceTemplateDom);
+                        detachedEntityDomAndPathList.emplace_back(detachedEntityDom, entityAliasPath);
+                        oldEntityAliases.emplace(entityId, commonRootEntityOwningInstance->get().GetEntityAlias(entityId)->get());
 
+                        commonRootEntityOwningInstance->get().DetachEntity(entityId).release();
+                    }
+                    
+                    PrefabUndoHelpers::RemoveEntities(detachedEntityDomAndPathList, commonRootInstanceTemplateId, undoBatch.GetUndoBatch());
+                }
+
+                // Detach the retrieved nested instances.
                 // When we create a prefab with other prefab instances, we have to remove the existing links between the source and 
                 // target templates of the other instances.
-                for (auto& nestedInstance : instances)
+                for (auto& instance : detachedInstances)
                 {
-                    AZStd::unique_ptr<Instance> outInstance = commonRootEntityOwningInstance->get().DetachNestedInstance(nestedInstance->GetInstanceAlias());
+                    const InstanceAlias instanceAlias = instance->GetInstanceAlias();
+
+                    AZStd::unique_ptr<Instance> detachedInstance =
+                        commonRootEntityOwningInstance->get().DetachNestedInstance(instanceAlias);
 
-                    LinkId detachingInstanceLinkId = nestedInstance->GetLinkId();
+                    LinkId detachingInstanceLinkId = detachedInstance->GetLinkId();
                     auto linkRef = m_prefabSystemComponentInterface->FindLink(detachingInstanceLinkId);
                     AZ_Assert(linkRef.has_value(), "Unable to find link with id '%llu' during prefab creation.", detachingInstanceLinkId);
 
                     PrefabDom linkPatches;
                     linkRef->get().GetLinkPatches(linkPatches, linkPatches.GetAllocator());
-                    nestedInstanceLinkPatchesMap.emplace(nestedInstance, AZStd::move(linkPatches));
+                    nestedInstanceLinkPatchesMap.emplace(detachedInstance.get(), AZStd::move(linkPatches));
 
-                    RemoveLink(outInstance, commonRootEntityOwningInstance->get().GetTemplateId(), undoBatch.GetUndoBatch());
+                    RemoveLink(detachedInstance, commonRootInstanceTemplateId, undoBatch.GetUndoBatch());
 
-                    instancePtrs.emplace_back(AZStd::move(outInstance));
+                    detachedInstancePtrs.emplace_back(AZStd::move(detachedInstance));
                 }
 
+                // Create new prefab instance.
                 auto prefabEditorEntityOwnershipInterface = AZ::Interface<PrefabEditorEntityOwnershipInterface>::Get();
                 if (!prefabEditorEntityOwnershipInterface)
                 {
@@ -192,9 +215,8 @@ namespace AzToolsFramework
                         "(PrefabEditorEntityOwnershipInterface unavailable)."));
                 }
 
-                // Create the Prefab
                 instanceToCreate = prefabEditorEntityOwnershipInterface->CreatePrefab(
-                    entities, AZStd::move(instancePtrs), m_prefabLoaderInterface->GenerateRelativePath(filePath),
+                    detachedEntities, AZStd::move(detachedInstancePtrs), m_prefabLoaderInterface->GenerateRelativePath(filePath),
                     commonRootEntityOwningInstance);
 
                 if (!instanceToCreate)
@@ -203,35 +225,53 @@ namespace AzToolsFramework
                         "(A null instance is returned)."));
                 }
 
-                containerEntityId = instanceToCreate->get().GetContainerEntityId();
+                // Note that changes to the new template DOM do not require undo/redo support. In other words, an undo after this operation
+                // will not change the newly created template DOM. This helps keep in-memory template in sync with in-disk prefab.
+
+                newContainerEntityId = instanceToCreate->get().GetContainerEntityId();
+                AZ::Entity* newContainerEntity = GetEntityById(newContainerEntityId);
+                const TemplateId newInstanceTemplateId = instanceToCreate->get().GetTemplateId();
+                PrefabDom& newInstanceTemplateDom = m_prefabSystemComponentInterface->FindTemplateDom(newInstanceTemplateId);
+
+                // Update and patch the container entity DOM in template with new components.
+                PrefabDom newContainerEntityDomWithComponents(&(newInstanceTemplateDom.GetAllocator()));
+                m_instanceToTemplateInterface->GenerateDomForEntity(newContainerEntityDomWithComponents, *newContainerEntity);
+                PrefabDomPath newContainerEntityAliasDomPath(m_instanceToTemplateInterface->GenerateEntityAliasPath(newContainerEntityId).c_str());
+                newContainerEntityAliasDomPath.Set(newInstanceTemplateDom, newContainerEntityDomWithComponents.Move());
 
                 // Apply the correct transform to the container for the new instance, and store the patch for use when creating the link.
-                PrefabDom patch = ApplyContainerTransformAndGeneratePatch(containerEntityId, commonRootEntityId,
+                PrefabDom linkPatches = ApplyContainerTransformAndGeneratePatch(newContainerEntityId, commonRootEntityId,
                     containerEntityTranslation, containerEntityRotation);
 
-                // Update the cache - this prevents these changes from being stored in the regular undo/redo nodes
-                m_prefabUndoCache.UpdateCache(containerEntityId);
+                // Capture the container entity DOM with parent data in transform.
+                // This DOM will be used later to generate patches for child sort update and they should not contain transform data differences.
+                PrefabDom newContainerEntityDomInitialStateWithTransform;
+                m_instanceToTemplateInterface->GenerateDomForEntity(newContainerEntityDomInitialStateWithTransform, *newContainerEntity);
 
-                // Parent the non-container top level entities to the container entity.
-                // Parenting the top level container entities will be done during the creation of links.
+                // Parent the non-container top level entities to the new container entity.
+                // It also patches the entity DOMs to the new template without undo/redo support.
                 EditorEntityContextNotificationBus::Broadcast(&EditorEntityContextNotification::SetForceAddEntitiesToBackFlag, true);
                 for (AZ::Entity* topLevelEntity : topLevelEntities)
                 {
-                    if (!IsInstanceContainerEntity(topLevelEntity->GetId()))
+                    const AZ::EntityId topLevelEntityId = topLevelEntity->GetId();
+
+                    // Note: Parenting the top level container entities will be done during the creation of links.
+                    if (!IsInstanceContainerEntity(topLevelEntityId))
                     {
-                        AZ::TransformBus::Event(topLevelEntity->GetId(), &AZ::TransformBus::Events::SetParent, containerEntityId);
+                        AZ::TransformBus::Event(topLevelEntityId, &AZ::TransformBus::Events::SetParent, newContainerEntityId);
+
+                        // Update the entity DOM with the latest transform data.
+                        PrefabDom topLevelContainerEntityDomWithParentInTransform(&(newInstanceTemplateDom.GetAllocator()));
+                        m_instanceToTemplateInterface->GenerateDomForEntity(
+                            topLevelContainerEntityDomWithParentInTransform, *topLevelEntity);
+
+                        PrefabDomPath entityAliasDomPath(m_instanceToTemplateInterface->GenerateEntityAliasPath(topLevelEntityId).c_str());
+                        entityAliasDomPath.Set(newInstanceTemplateDom, topLevelContainerEntityDomWithParentInTransform.Move());
                     }
                 }
                 EditorEntityContextNotificationBus::Broadcast(&EditorEntityContextNotification::SetForceAddEntitiesToBackFlag, false);
 
-                // Update the template of the instance since the entities are modified since the template creation.
-                Prefab::PrefabDom serializedInstance;
-                if (m_instanceToTemplateInterface->GenerateDomForInstance(serializedInstance, instanceToCreate->get()))
-                {
-                    m_prefabSystemComponentInterface->UpdatePrefabTemplate(instanceToCreate->get().GetTemplateId(), serializedInstance);
-                }
-
-                // Set up links between the created prefab instance and its nested instances.
+                // Set up links between the created prefab instance and its nested instances without undo/redo support.
                 instanceToCreate->get().GetNestedInstances([&](AZStd::unique_ptr<Instance>& nestedInstance) {
                     AZ_Assert(nestedInstance, "Invalid nested instance found in the new prefab created.");
 
@@ -251,21 +291,21 @@ namespace AzToolsFramework
 
                     // These link creations shouldn't be undone because that would put the template in a non-usable state if a user
                     // chooses to instantiate the template after undoing the creation.
-                    CreateLink(*nestedInstance, instanceToCreate->get().GetTemplateId(), undoBatch.GetUndoBatch(), AZStd::move(previousPatch), false);
+                    CreateLink(*nestedInstance, newInstanceTemplateId, undoBatch.GetUndoBatch(), AZStd::move(previousPatch), false);
 
                     // If this nested instance's container is a top level entity in the new prefab, re-parent it and apply the change.
                     if (AZStd::find(topLevelEntities.begin(), topLevelEntities.end(), &nestedInstanceContainerEntity->get()) != topLevelEntities.end())
                     {
-                        Prefab::PrefabDom containerEntityDomBefore;
-                        m_instanceToTemplateInterface->GenerateDomForEntity(containerEntityDomBefore, *nestedInstanceContainerEntity);
+                        PrefabDom nestedContainerEntityDomBefore;
+                        m_instanceToTemplateInterface->GenerateDomForEntity(nestedContainerEntityDomBefore, *nestedInstanceContainerEntity);
 
-                        AZ::TransformBus::Event(nestedInstanceContainerEntityId, &AZ::TransformBus::Events::SetParent, containerEntityId);
+                        AZ::TransformBus::Event(nestedInstanceContainerEntityId, &AZ::TransformBus::Events::SetParent, newContainerEntityId);
 
-                        PrefabDom containerEntityDomAfter;
-                        m_instanceToTemplateInterface->GenerateDomForEntity(containerEntityDomAfter, *nestedInstanceContainerEntity);
+                        PrefabDom nestedContainerEntityDomAfter;
+                        m_instanceToTemplateInterface->GenerateDomForEntity(nestedContainerEntityDomAfter, *nestedInstanceContainerEntity);
 
                         PrefabDom reparentPatch;
-                        m_instanceToTemplateInterface->GeneratePatch(reparentPatch, containerEntityDomBefore, containerEntityDomAfter);
+                        m_instanceToTemplateInterface->GeneratePatch(reparentPatch, nestedContainerEntityDomBefore, nestedContainerEntityDomAfter);
                         m_instanceToTemplateInterface->AppendEntityAliasToPatchPaths(reparentPatch, nestedInstanceContainerEntityId);
 
                         // We won't parent this undo node to the undo batch so that the newly created template and link will remain
@@ -276,33 +316,38 @@ namespace AzToolsFramework
                     }
                 });
 
-                // Create a link between the templates of the newly created instance and the instance it's being parented under.
-                CreateLink(
-                    instanceToCreate->get(), commonRootEntityOwningInstance->get().GetTemplateId(), undoBatch.GetUndoBatch(),
-                    AZStd::move(patch));
+                // Patch the new template with child sort array update without undo/redo support.
+                PrefabDom newContainerEntityDomWithTransformAndChildSortOrder;
+                m_instanceToTemplateInterface->GenerateDomForEntity(newContainerEntityDomWithTransformAndChildSortOrder, *newContainerEntity);
 
-                // Reset the transform of the container entity so that the new values aren't saved in the new prefab's dom.
-                // The new values were saved in the link, so propagation will apply them correctly.
-                {
-                    AZ::Entity* containerEntity = GetEntityById(containerEntityId);
-
-                    PrefabDom containerBeforeReset;
-                    m_instanceToTemplateInterface->GenerateDomForEntity(containerBeforeReset, *containerEntity);
-
-                    AZ::TransformBus::Event(containerEntityId, &AZ::TransformBus::Events::SetParent, AZ::EntityId());
-                    AZ::TransformBus::Event(containerEntityId, &AZ::TransformBus::Events::SetLocalTM, AZ::Transform::CreateIdentity());
+                PrefabDom childSortOrderUpdatePatches;
+                m_instanceToTemplateInterface->GeneratePatch(childSortOrderUpdatePatches,
+                    newContainerEntityDomInitialStateWithTransform, newContainerEntityDomWithTransformAndChildSortOrder);
+                m_instanceToTemplateInterface->AppendEntityAliasToPatchPaths(childSortOrderUpdatePatches, newContainerEntityId);
+                m_instanceToTemplateInterface->PatchTemplate(childSortOrderUpdatePatches, newInstanceTemplateId);
 
-                    PrefabDom containerAfterReset;
-                    m_instanceToTemplateInterface->GenerateDomForEntity(containerAfterReset, *containerEntity);
+                // Create a link between the new prefab template and its parent template with undo/redo support.
+                CreateLink(instanceToCreate->get(), commonRootInstanceTemplateId, undoBatch.GetUndoBatch(), AZStd::move(linkPatches));
 
-                    // Update the state of the entity
-                    auto templateId = instanceToCreate->get().GetTemplateId();
-
-                    PrefabDom transformPatch;
-                    m_instanceToTemplateInterface->GeneratePatch(transformPatch, containerBeforeReset, containerAfterReset);
-                    m_instanceToTemplateInterface->AppendEntityAliasToPatchPaths(transformPatch, containerEntityId);
+                // Update common parent entity DOM with the new prefab instance as child with undo/redo support.
+                if (commonRootEntityId.IsValid())
+                {
+                    PrefabDom commonRootEntityDomAfterAddingNewPrefab;
+                    m_instanceToTemplateInterface->GenerateDomForEntity(commonRootEntityDomAfterAddingNewPrefab,
+                        *GetEntityById(commonRootEntityId));
+
+                    PrefabUndoHelpers::UpdateEntity(
+                        commonRootEntityDomBeforeDetaching,
+                        commonRootEntityDomAfterAddingNewPrefab,
+                        commonRootEntityId,
+                        undoBatch.GetUndoBatch());
+                }
 
-                    m_instanceToTemplateInterface->PatchTemplate(transformPatch, templateId);
+                // Update undo cache.
+                m_prefabUndoCache.UpdateCache(newContainerEntityId);
+                for (AZ::Entity* topLevelEntity : topLevelEntities)
+                {
+                    m_prefabUndoCache.UpdateCache(topLevelEntity->GetId());
                 }
 
                 // This clears any entities marked as dirty due to reparenting of entities during the process of creating a prefab.
@@ -312,14 +357,14 @@ namespace AzToolsFramework
 
                 // Select Container Entity
                 {
-                    const EntityIdList selectedEntities = EntityIdList{ containerEntityId };
+                    const EntityIdList selectedEntities = EntityIdList{ newContainerEntityId };
                     auto selectionUndo = aznew SelectionCommand(selectedEntities, "Select Prefab Container Entity");
                     selectionUndo->SetParent(undoBatch.GetUndoBatch());
                     ToolsApplicationRequestBus::Broadcast(&ToolsApplicationRequestBus::Events::SetSelectedEntities, selectedEntities);
                 }
             }
 
-            return AZ::Success(containerEntityId);
+            return AZ::Success(newContainerEntityId);
         }
 
         CreatePrefabResult PrefabPublicHandler::CreatePrefabInDisk(const EntityIdList& entityIds, AZ::IO::PathView filePath)
@@ -335,7 +380,7 @@ namespace AzToolsFramework
                 if (!m_prefabLoaderInterface->SaveTemplateToFile(templateId, filePath))
                 {
                     AZStd::string_view filePathString(filePath);
-                    return AZ::Failure(AZStd::string::format(
+                    return AZ::Failure(AZStd::string::format("CreatePrefabInDisk - "
                         "Could not save the newly created prefab to file path %.*s - internal error ",
                         AZ_STRING_ARG(filePathString)));
                 }

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabPublicHandler.h

@@ -43,10 +43,8 @@ namespace AzToolsFramework
             void UnregisterPrefabPublicHandlerInterface();
 
             // PrefabPublicInterface...
-            CreatePrefabResult CreatePrefabInDisk(
-                const EntityIdList& entityIds, AZ::IO::PathView filePath) override;
-            CreatePrefabResult CreatePrefabInMemory(
-                const EntityIdList& entityIds, AZ::IO::PathView filePath) override;
+            CreatePrefabResult CreatePrefabInDisk(const EntityIdList& entityIds, AZ::IO::PathView filePath) override;
+            CreatePrefabResult CreatePrefabInMemory(const EntityIdList& entityIds, AZ::IO::PathView filePath) override;
             InstantiatePrefabResult InstantiatePrefab(
                 AZStd::string_view filePath, AZ::EntityId parentId, const AZ::Vector3& position) override;
             PrefabOperationResult SavePrefab(AZ::IO::Path filePath) override;

Code/Framework/AzToolsFramework/AzToolsFramework/Prefab/PrefabPublicInterface.h

@@ -44,24 +44,24 @@ namespace AzToolsFramework
             /**
              * Create a prefab out of the entities provided, at the path provided, and save it in disk immediately.
              * Automatically detects descendants of entities, and discerns between entities and child instances.
+             * Note: It calls CreatePrefabInMemory internally and then saves the new template in disk.
              * @param entityIds The entities that should form the new prefab (along with their descendants).
              * @param filePath The absolute path for the new prefab file.
              * @return An outcome object with an entityId of the new prefab's container entity;
              *  on failure, it comes with an error message detailing the cause of the error.
              */
-            virtual CreatePrefabResult CreatePrefabInDisk(
-                const EntityIdList& entityIds, AZ::IO::PathView filePath) = 0;
+            virtual CreatePrefabResult CreatePrefabInDisk(const EntityIdList& entityIds, AZ::IO::PathView filePath) = 0;
 
             /**
              * Create a prefab out of the entities provided, at the path provided, and keep it in memory.
              * Automatically detects descendants of entities, and discerns between entities and child instances.
+             * Note: The newly created prefab template cannot be undo/redo. Undo will not remove the templace in system.
              * @param entityIds The entities that should form the new prefab (along with their descendants).
              * @param filePath The absolute path for the new prefab file.
              * @return An outcome object with an entityId of the new prefab's container entity;
              *  on failure, it comes with an error message detailing the cause of the error.
              */
-            virtual CreatePrefabResult CreatePrefabInMemory(
-                const EntityIdList& entityIds, AZ::IO::PathView filePath) = 0;
+            virtual CreatePrefabResult CreatePrefabInMemory(const EntityIdList& entityIds, AZ::IO::PathView filePath) = 0;
 
             /**
              * Instantiate a prefab from a prefab file.

Code/Framework/AzToolsFramework/Tests/Prefab/PrefabDeleteTests.cpp

@@ -37,10 +37,8 @@ namespace UnitTest
 
     TEST_F(PrefabDeleteTest, DeleteEntitiesAndAllDescendantsInInstance_DeleteSinglePrefabSucceeds)
     {
-        PrefabEntityResult createEntityResult = m_prefabPublicInterface->CreateEntity(AZ::EntityId(), AZ::Vector3());
-
         // Verify that a valid entity is created.
-        AZ::EntityId createdEntityId = createEntityResult.GetValue();
+        AZ::EntityId createdEntityId = CreateEntityUnderRootPrefab("EntityUnderRootPrefab");
         ASSERT_TRUE(createdEntityId.IsValid());
         AZ::Entity* createdEntity = AzToolsFramework::GetEntityById(createdEntityId);
         ASSERT_TRUE(createdEntity != nullptr);
@@ -101,10 +99,8 @@ namespace UnitTest
 
     TEST_F(PrefabDeleteTest, DeleteEntitiesAndAllDescendantsInInstance_DeletingEntityDeletesChildPrefabToo)
     {
-        PrefabEntityResult entityToBePutUnderPrefabResult = m_prefabPublicInterface->CreateEntity(AZ::EntityId(), AZ::Vector3());
-
         // Verify that a valid entity is created that will be put in a prefab later.
-        AZ::EntityId entityToBePutUnderPrefabId = entityToBePutUnderPrefabResult.GetValue();
+        AZ::EntityId entityToBePutUnderPrefabId = CreateEntityUnderRootPrefab("EntityToBePutUnderPrefab");
         ASSERT_TRUE(entityToBePutUnderPrefabId.IsValid());
         AZ::Entity* entityToBePutUnderPrefab = AzToolsFramework::GetEntityById(entityToBePutUnderPrefabId);
         ASSERT_TRUE(entityToBePutUnderPrefab != nullptr);