NetBindComponent::SetOwningConnectionId is for a hierarchy based on its root's network entity

- for hierarchies, on both servers and clients NetBindComponent::SetOwningConnectionId is set for the all involved entities based on the root entity owning connection id
- added new unit tests to cover these scenarios
- all unit tests pass, hierarchy benchmarks remain unaffected

Gems/Multiplayer/Code/Include/Multiplayer/Components/NetworkHierarchyChildComponent.h

@@ -58,11 +58,10 @@ namespace Multiplayer
         void BindNetworkHierarchyLeaveEventHandler(NetworkHierarchyLeaveEvent::Handler& handler) override;
         //! @}
 
-    protected:
-        //! Used by @NetworkHierarchyRootComponent
-        void SetTopLevelHierarchyRootEntity(AZ::Entity* hierarchyRoot);
-
     private:
+        //! Used by @NetworkHierarchyRootComponent
+        void SetTopLevelHierarchyRootEntity(AZ::Entity* previousHierarchyRoot, AZ::Entity* newHierarchyRoot);
+
         AZ::ChildChangedEvent::Handler m_childChangedHandler;
 
         void OnChildChanged(AZ::ChildChangeType type, AZ::EntityId child);
@@ -80,5 +79,8 @@ namespace Multiplayer
         bool m_isHierarchyEnabled = true;
 
         void NotifyChildrenHierarchyDisbanded();
+
+        AzNetworking::ConnectionId m_previousOwningConnectionId = AzNetworking::InvalidConnectionId;
+        void SetOwningConnectionId(AzNetworking::ConnectionId connectionId) override;
     };
 }

Gems/Multiplayer/Code/Include/Multiplayer/Components/NetworkHierarchyRootComponent.h

@@ -60,10 +60,9 @@ namespace Multiplayer
 
         bool SerializeEntityCorrection(AzNetworking::ISerializer& serializer);
 
-    protected:
-        void SetTopLevelHierarchyRootEntity(AZ::Entity* hierarchyRoot);
-
     private:
+        void SetTopLevelHierarchyRootEntity(AZ::Entity* previousHierarchyRoot, AZ::Entity* newHierarchyRoot);
+
         AZ::ChildChangedEvent::Handler m_childChangedHandler;
         AZ::ParentChangedEvent::Handler m_parentChangedHandler;
 
@@ -80,16 +79,19 @@ namespace Multiplayer
 
         //! Rebuilds hierarchy starting from this root component's entity.
         void RebuildHierarchy();
-        
+
         //! @param underEntity Walk the child entities that belong to @underEntity and consider adding them to the hierarchy.
         //! Builds the hierarchy using breadth-first iterative method.
         void InternalBuildHierarchyList(AZ::Entity* underEntity);
 
-        void SetRootForEntity(AZ::Entity* root, const AZ::Entity* childEntity);
+        void SetRootForEntity(AZ::Entity* previousKnownRoot, AZ::Entity* newRoot, const AZ::Entity* childEntity);
 
         //! Set to false when deactivating or otherwise not to be included in hierarchy considerations.
         bool m_isHierarchyEnabled = true;
 
+        AzNetworking::ConnectionId m_previousOwningConnectionId = AzNetworking::InvalidConnectionId;
+        void SetOwningConnectionId(AzNetworking::ConnectionId connectionId) override;
+
         friend class HierarchyBenchmarkBase;
     };
 

Gems/Multiplayer/Code/Source/Components/NetworkHierarchyChildComponent.cpp

@@ -129,34 +129,48 @@ namespace Multiplayer
         handler.Connect(m_networkHierarchyLeaveEvent);
     }
 
-    void NetworkHierarchyChildComponent::SetTopLevelHierarchyRootEntity(AZ::Entity* hierarchyRoot)
+    void NetworkHierarchyChildComponent::SetTopLevelHierarchyRootEntity(AZ::Entity* previousHierarchyRoot, AZ::Entity* newHierarchyRoot)
     {
-        if (m_rootEntity != hierarchyRoot)
+        if (newHierarchyRoot)
         {
-            m_rootEntity = hierarchyRoot;
+            if (m_rootEntity != newHierarchyRoot)
+            {
+                m_rootEntity = newHierarchyRoot;
+
+                if (HasController() && GetNetBindComponent()->GetNetEntityRole() == NetEntityRole::Authority)
+                {
+                    NetworkHierarchyChildComponentController* controller = static_cast<NetworkHierarchyChildComponentController*>(GetController());
+                    const NetEntityId netRootId = GetNetworkEntityManager()->GetNetEntityIdById(m_rootEntity->GetId());
+                    controller->SetHierarchyRoot(netRootId);
+                }
+
+                GetNetBindComponent()->SetOwningConnectionId(m_rootEntity->FindComponent<NetBindComponent>()->GetOwningConnectionId());
+                m_networkHierarchyChangedEvent.Signal(m_rootEntity->GetId());
+            }
+        }
+        else if ((previousHierarchyRoot && m_rootEntity == previousHierarchyRoot) || !previousHierarchyRoot)
+        {
+            m_rootEntity = nullptr;
 
             if (HasController() && GetNetBindComponent()->GetNetEntityRole() == NetEntityRole::Authority)
             {
                 NetworkHierarchyChildComponentController* controller = static_cast<NetworkHierarchyChildComponentController*>(GetController());
-                if (m_rootEntity)
-                {
-                    const NetEntityId netRootId = GetNetworkEntityManager()->GetNetEntityIdById(m_rootEntity->GetId());
-                    controller->SetHierarchyRoot(netRootId);
-
-                    m_networkHierarchyChangedEvent.Signal(m_rootEntity->GetId());
-                }
-                else
-                {
-                    controller->SetHierarchyRoot(InvalidNetEntityId);
-
-                    m_networkHierarchyLeaveEvent.Signal();
-                }
+                controller->SetHierarchyRoot(InvalidNetEntityId);
             }
 
-            if (m_rootEntity == nullptr)
-            {
-                NotifyChildrenHierarchyDisbanded();
-            }
+            GetNetBindComponent()->SetOwningConnectionId(m_previousOwningConnectionId);
+            m_networkHierarchyLeaveEvent.Signal();
+
+            NotifyChildrenHierarchyDisbanded();
+        }
+    }
+
+    void NetworkHierarchyChildComponent::SetOwningConnectionId(AzNetworking::ConnectionId connectionId)
+    {
+        NetworkHierarchyChildComponentBase::SetOwningConnectionId(connectionId);
+        if (IsHierarchicalChild() == false)
+        {
+            m_previousOwningConnectionId = connectionId;
         }
     }
 
@@ -180,14 +194,18 @@ namespace Multiplayer
             if (m_rootEntity != newRoot)
             {
                 m_rootEntity = newRoot;
+
+                m_previousOwningConnectionId = GetNetBindComponent()->GetOwningConnectionId();
+                GetNetBindComponent()->SetOwningConnectionId(m_rootEntity->FindComponent<NetBindComponent>()->GetOwningConnectionId());
+
                 m_networkHierarchyChangedEvent.Signal(m_rootEntity->GetId());
             }
         }
         else
         {
+            GetNetBindComponent()->SetOwningConnectionId(m_previousOwningConnectionId);
             m_isHierarchyEnabled = false;
             m_rootEntity = nullptr;
-            m_networkHierarchyLeaveEvent.Signal();
         }
     }
 
@@ -203,11 +221,11 @@ namespace Multiplayer
             {
                 if (auto* hierarchyChildComponent = childEntity->FindComponent<NetworkHierarchyChildComponent>())
                 {
-                    hierarchyChildComponent->SetTopLevelHierarchyRootEntity(nullptr);
+                    hierarchyChildComponent->SetTopLevelHierarchyRootEntity(nullptr, nullptr);
                 }
                 else if (auto* hierarchyRootComponent = childEntity->FindComponent<NetworkHierarchyRootComponent>())
                 {
-                    hierarchyRootComponent->SetTopLevelHierarchyRootEntity(nullptr);
+                    hierarchyRootComponent->SetTopLevelHierarchyRootEntity(nullptr, nullptr);
                 }
             }
         }

Gems/Multiplayer/Code/Source/Components/NetworkHierarchyRootComponent.cpp

@@ -107,7 +107,7 @@ namespace Multiplayer
             {
                 if (const AZ::Entity* childEntity = AZ::Interface<AZ::ComponentApplicationRequests>::Get()->FindEntity(childEntityId))
                 {
-                    SetRootForEntity(nullptr, childEntity);
+                    SetRootForEntity(GetEntity(), nullptr, childEntity);
                 }
             }
         }
@@ -209,7 +209,7 @@ namespace Multiplayer
             auto [rootComponent, childComponent] = GetHierarchyComponents(parentEntity);
             if (rootComponent == nullptr && childComponent == nullptr)
             {
-                RebuildHierarchy();
+                SetRootForEntity(nullptr, nullptr, GetEntity());
             }
             else
             {
@@ -219,7 +219,7 @@ namespace Multiplayer
         else
         {
             // Detached from parent
-            RebuildHierarchy();
+            SetRootForEntity(nullptr, nullptr, GetEntity());
         }
     }
 
@@ -247,14 +247,14 @@ namespace Multiplayer
             {
                 // This is a newly added entity to the network hierarchy.
                 hierarchyChanged = true;
-                SetRootForEntity(GetEntity(), currentEntity);
+                SetRootForEntity(nullptr, GetEntity(), currentEntity);
             }
         }
 
         // These entities were removed since last rebuild.
         for (const AZ::Entity* previousEntity : previousEntities)
         {
-            SetRootForEntity(nullptr, previousEntity);
+            SetRootForEntity(GetEntity(), nullptr, previousEntity);
         }
 
         if (!previousEntities.empty())
@@ -307,45 +307,66 @@ namespace Multiplayer
         }
     }
 
-    void NetworkHierarchyRootComponent::SetRootForEntity(AZ::Entity* root, const AZ::Entity* childEntity)
+    void NetworkHierarchyRootComponent::SetRootForEntity(AZ::Entity* previousKnownRoot, AZ::Entity* newRoot, const AZ::Entity* childEntity)
     {
         auto [hierarchyRootComponent, hierarchyChildComponent] = GetHierarchyComponents(childEntity);
 
         if (hierarchyChildComponent)
         {
-            hierarchyChildComponent->SetTopLevelHierarchyRootEntity(root);
+            hierarchyChildComponent->SetTopLevelHierarchyRootEntity(previousKnownRoot, newRoot);
         }
         else if (hierarchyRootComponent)
         {
-            hierarchyRootComponent->SetTopLevelHierarchyRootEntity(root);
+            hierarchyRootComponent->SetTopLevelHierarchyRootEntity(previousKnownRoot, newRoot);
         }
     }
 
-    void NetworkHierarchyRootComponent::SetTopLevelHierarchyRootEntity(AZ::Entity* hierarchyRoot)
+    void NetworkHierarchyRootComponent::SetTopLevelHierarchyRootEntity(AZ::Entity* previousHierarchyRoot, AZ::Entity* newHierarchyRoot)
     {
-        m_rootEntity = hierarchyRoot;
-
-        if (HasController() && GetNetBindComponent()->GetNetEntityRole() == NetEntityRole::Authority)
+        if (newHierarchyRoot)
         {
-            NetworkHierarchyChildComponentController* controller = static_cast<NetworkHierarchyChildComponentController*>(GetController());
-            if (hierarchyRoot)
+            if (m_rootEntity != newHierarchyRoot)
             {
-                const NetEntityId netRootId = GetNetworkEntityManager()->GetNetEntityIdById(hierarchyRoot->GetId());
-                controller->SetHierarchyRoot(netRootId);
-            }
-            else
-            {
-                controller->SetHierarchyRoot(InvalidNetEntityId);
+                m_rootEntity = newHierarchyRoot;
+
+                if (HasController() && GetNetBindComponent()->GetNetEntityRole() == NetEntityRole::Authority)
+                {
+                    NetworkHierarchyRootComponentController* controller = static_cast<NetworkHierarchyRootComponentController*>(GetController());
+                    const NetEntityId netRootId = GetNetworkEntityManager()->GetNetEntityIdById(m_rootEntity->GetId());
+                    controller->SetHierarchyRoot(netRootId);
+                }
+
+                GetNetBindComponent()->SetOwningConnectionId(m_rootEntity->FindComponent<NetBindComponent>()->GetOwningConnectionId());
+                m_networkHierarchyChangedEvent.Signal(m_rootEntity->GetId());
             }
         }
-
-        if (m_rootEntity == nullptr)
+        else if ((previousHierarchyRoot && m_rootEntity == previousHierarchyRoot) || !previousHierarchyRoot)
         {
+            m_rootEntity = nullptr;
+
+            if (HasController() && GetNetBindComponent()->GetNetEntityRole() == NetEntityRole::Authority)
+            {
+                NetworkHierarchyRootComponentController* controller = static_cast<NetworkHierarchyRootComponentController*>(GetController());
+                controller->SetHierarchyRoot(InvalidNetEntityId);
+            }
+
+            GetNetBindComponent()->SetOwningConnectionId(m_previousOwningConnectionId);
+            m_networkHierarchyLeaveEvent.Signal();
+
             // We lost the parent hierarchical entity, so as a root we need to re-build our own hierarchy.
             RebuildHierarchy();
         }
     }
 
+    void NetworkHierarchyRootComponent::SetOwningConnectionId(AzNetworking::ConnectionId connectionId)
+    {
+        NetworkHierarchyRootComponentBase::SetOwningConnectionId(connectionId);
+        if (IsHierarchicalChild() == false)
+        {
+            m_previousOwningConnectionId = connectionId;
+        }
+    }
+
     NetworkHierarchyRootComponentController::NetworkHierarchyRootComponentController(NetworkHierarchyRootComponent& parent)
         : NetworkHierarchyRootComponentControllerBase(parent)
     {
@@ -370,7 +391,7 @@ namespace Multiplayer
     void NetworkHierarchyRootComponentController::CreateInput(Multiplayer::NetworkInput& input, float deltaTime)
     {
         NetworkHierarchyRootComponent& component = GetParent();
-        if(!component.IsHierarchicalRoot())
+        if (!component.IsHierarchicalRoot())
         {
             return;
         }
@@ -386,7 +407,7 @@ namespace Multiplayer
 
         for (AZ::Entity* child : entities)
         {
-            if(child == component.GetEntity())
+            if (child == component.GetEntity())
             {
                 continue; // Avoid infinite recursion
             }

Gems/Multiplayer/Code/Tests/ClientHierarchyTests.cpp

@@ -302,6 +302,36 @@ namespace Multiplayer
         SetHierarchyRootFieldOnNetworkHierarchyChildOnClient(m_child->m_entity, InvalidNetEntityId);
     }
 
+    TEST_F(ClientSimpleHierarchyTests, ChildHasOwningConnectionIdOfParent)
+    {
+        // disconnect and assign new connection ids
+        SetParentIdOnNetworkTransform(m_child->m_entity, InvalidNetEntityId);
+        SetHierarchyRootFieldOnNetworkHierarchyChildOnClient(m_child->m_entity, InvalidNetEntityId);
+
+        m_root->m_entity->FindComponent<NetBindComponent>()->SetOwningConnectionId(ConnectionId{ 1 });
+        m_child->m_entity->FindComponent<NetBindComponent>()->SetOwningConnectionId(ConnectionId{ 2 });
+
+        const ConnectionId previousConnectionId = m_child->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId();
+
+        // re-attach, child's owning connection id should then be root's connection id
+        SetParentIdOnNetworkTransform(m_child->m_entity, RootNetEntityId);
+        SetHierarchyRootFieldOnNetworkHierarchyChildOnClient(m_child->m_entity, RootNetEntityId);
+
+        EXPECT_EQ(
+            m_child->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            m_root->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId()
+        );
+
+        // detach, the child should roll back to his previous owning connection id
+        SetParentIdOnNetworkTransform(m_child->m_entity, InvalidNetEntityId);
+        SetHierarchyRootFieldOnNetworkHierarchyChildOnClient(m_child->m_entity, InvalidNetEntityId);
+
+        EXPECT_EQ(
+            m_child->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            previousConnectionId
+        );
+    }
+
     /*
      * Parent -> Child -> ChildOfChild
      */
@@ -396,7 +426,7 @@ namespace Multiplayer
         using MultiplayerTest::TestMultiplayerComponentNetworkInput;
 
         auto* rootNetBind = m_root->m_entity->FindComponent<NetBindComponent>();
-        
+
         NetworkInputArray inputArray(rootNetBind->GetEntityHandle());
         NetworkInput& input = inputArray[0];
 

Gems/Multiplayer/Code/Tests/ServerHierarchyTests.cpp

@@ -195,6 +195,49 @@ namespace Multiplayer
         m_child->m_entity.reset();
     }
 
+    TEST_F(ServerSimpleHierarchyTests, ChildPointsToRootAfterReattachment)
+    {
+        m_child->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(AZ::EntityId());
+
+        EXPECT_EQ(
+            m_child->m_entity->FindComponent<NetworkHierarchyChildComponent>()->GetHierarchyRoot(),
+            InvalidNetEntityId
+        );
+
+        m_child->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(m_root->m_entity->GetId());
+
+        EXPECT_EQ(
+            m_child->m_entity->FindComponent<NetworkHierarchyChildComponent>()->GetHierarchyRoot(),
+            m_root->m_entity->FindComponent<NetBindComponent>()->GetNetEntityId()
+        );
+    }
+
+    TEST_F(ServerSimpleHierarchyTests, ChildHasOwningConnectionIdOfParent)
+    {
+        // disconnect and assign new connection ids
+        m_child->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(AZ::EntityId());
+        m_root->m_entity->FindComponent<NetBindComponent>()->SetOwningConnectionId(ConnectionId{ 1 });
+        m_child->m_entity->FindComponent<NetBindComponent>()->SetOwningConnectionId(ConnectionId{ 2 });
+
+        const ConnectionId previousConnectionId = m_child->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId();
+
+        // re-attach, child's owning connection id should then be root's connection id
+        m_child->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(m_root->m_entity->GetId());
+
+        EXPECT_EQ(
+            m_child->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            m_root->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId()
+        );
+
+        // detach, the child should roll back to his previous owning connection id
+        m_child->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(AZ::EntityId());
+
+        EXPECT_EQ(
+            m_child->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            previousConnectionId
+        );
+    }
+
     /*
      * Parent -> Child -> ChildOfChild
      */
@@ -410,7 +453,7 @@ namespace Multiplayer
     {
         MockNetworkHierarchyCallbackHandler mock;
         EXPECT_CALL(mock, OnNetworkHierarchyUpdated(m_root->m_entity->GetId())).Times(2);
-        
+
         m_root->m_entity->FindComponent<NetworkHierarchyRootComponent>()->BindNetworkHierarchyChangedEventHandler(mock.m_changedHandler);
 
         m_child->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(AZ::EntityId());
@@ -822,6 +865,22 @@ namespace Multiplayer
         }
     }
 
+    TEST_F(ServerHierarchyOfHierarchyTests, InnerChildrenPointToInnerRootAfterDetachmentFromTopRoot)
+    {
+        m_root2->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(m_root->m_entity->GetId());
+        // detach
+        m_root2->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(AZ::EntityId());
+
+        EXPECT_EQ(
+            m_child2->m_entity->FindComponent<NetworkHierarchyChildComponent>()->GetHierarchyRoot(),
+            m_root2->m_entity->FindComponent<NetBindComponent>()->GetNetEntityId()
+        );
+        EXPECT_EQ(
+            m_childOfChild2->m_entity->FindComponent<NetworkHierarchyChildComponent>()->GetHierarchyRoot(),
+            m_root2->m_entity->FindComponent<NetBindComponent>()->GetNetEntityId()
+        );
+    }
+
     TEST_F(ServerHierarchyOfHierarchyTests, Inner_Root_Has_Child_References_After_Detachment_From_Child_Of_Child)
     {
         m_root2->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(m_childOfChild->m_entity->GetId());
@@ -1005,6 +1064,59 @@ namespace Multiplayer
         m_console->GetCvarValue<uint32_t>("bg_hierarchyEntityMaxLimit", currentMaxLimit);
     }
 
+    TEST_F(ServerHierarchyOfHierarchyTests, InnerRootAndItsChildrenHaveOwningConnectionIdOfTopRoot)
+    {
+        // Assign new connection ids.
+        m_root->m_entity->FindComponent<NetBindComponent>()->SetOwningConnectionId(ConnectionId{ 1 });
+        m_root2->m_entity->FindComponent<NetBindComponent>()->SetOwningConnectionId(ConnectionId{ 2 });
+
+        // Attach then inner hierarchy's owning connection id should then be top root's connection id.
+        m_root2->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(m_childOfChild->m_entity->GetId());
+
+        EXPECT_EQ(
+            m_root2->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            m_root->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId()
+        );
+
+        EXPECT_EQ(
+            m_child2->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            m_root->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId()
+        );
+
+        EXPECT_EQ(
+            m_childOfChild2->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            m_root->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId()
+        );
+    }
+
+    TEST_F(ServerHierarchyOfHierarchyTests, InnerRootAndItsChildrenHaveTheirOriginalOwningConnectionIdAfterDetachingFromTopRoot)
+    {
+        // Assign new connection ids.
+        m_root->m_entity->FindComponent<NetBindComponent>()->SetOwningConnectionId(ConnectionId{ 1 });
+        m_root2->m_entity->FindComponent<NetBindComponent>()->SetOwningConnectionId(ConnectionId{ 2 });
+
+        // Attach then inner hierarchy's owning connection id should then be top root's connection id.
+        m_root2->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(m_childOfChild->m_entity->GetId());
+
+        // detach, inner hierarchy should roll back to his previous owning connection id
+        m_root2->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(AZ::EntityId());
+
+        EXPECT_EQ(
+            m_root2->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            ConnectionId{ 2 }
+        );
+
+        EXPECT_EQ(
+            m_child2->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            m_root2->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId()
+        );
+
+        EXPECT_EQ(
+            m_childOfChild2->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId(),
+            m_root2->m_entity->FindComponent<NetBindComponent>()->GetOwningConnectionId()
+        );
+    }
+
     /*
      * Parent -> Child -> ChildOfChild (not marked as in a hierarchy)
      */
@@ -1242,15 +1354,15 @@ namespace Multiplayer
         );
     }
 
-    TEST_F(ServerHierarchyWithThreeRoots, ReattachMiddleChildWhileLastChildGetsLeaveEventOnce)
+    TEST_F(ServerHierarchyWithThreeRoots, InnerRootLeftTopRootThenLastChildGetsJoinedEventOnce)
     {
         m_root2->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(m_childOfChild->m_entity->GetId());
         m_root3->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(m_childOfChild->m_entity->GetId());
 
         MockNetworkHierarchyCallbackHandler mock;
-        EXPECT_CALL(mock, OnNetworkHierarchyLeave());
-        
-        m_childOfChild3->m_entity->FindComponent<NetworkHierarchyChildComponent>()->BindNetworkHierarchyLeaveEventHandler(mock.m_leaveHandler);
+        EXPECT_CALL(mock, OnNetworkHierarchyUpdated(m_root3->m_entity->GetId()));
+
+        m_childOfChild3->m_entity->FindComponent<NetworkHierarchyChildComponent>()->BindNetworkHierarchyChangedEventHandler(mock.m_changedHandler);
 
         m_child->m_entity->FindComponent<AzFramework::TransformComponent>()->SetParent(AZ::EntityId());
     }