Add a test for resetting extended ownership flags

Author: EmandM

com.unity.netcode.gameobjects/Tests/Runtime/DistributedAuthority/OwnershipPermissionsTests.cs

@@ -499,6 +499,108 @@ private IEnumerator ValidateRequestAndPermissionsChangeRaceCondition(NetworkObje
             AssertOnTimeout($"[{newStatus}][Set RequestRequired][Permissions mismatch] {firstClient.name}");
         }
 
+        [UnityTest]
+        public IEnumerator ResetExtendedOwnershipFlagsOnRespawn()
+        {
+            var extendedFlags = (NetworkObject.OwnershipStatusExtended[])Enum.GetValues(typeof(NetworkObject.OwnershipStatusExtended));
+            Assert.That(extendedFlags.Length, Is.LessThan(NumberOfClients), "This test should have at least one more non-authority client than extended flags!");
+
+            var spawnedObjects = new List<NetworkObject>();
+            while (spawnedObjects.Count <= extendedFlags.Length)
+            {
+                var ownerClient = GetNonAuthorityNetworkManager(spawnedObjects.Count);
+                var spawnedObject = SpawnObject(m_PermissionsObject, ownerClient).GetComponent<NetworkObject>();
+                spawnedObject.SetOwnershipStatus(NetworkObject.OwnershipStatus.RequestRequired, true);
+                spawnedObjects.Add(spawnedObject);
+            }
+
+            yield return WaitForSpawnedOnAllOrTimeOut(spawnedObjects);
+            AssertOnTimeout("[InitialSpawn] Not all objects are spawned on all clients!");
+
+            for (int i = 0; i < extendedFlags.Length; i++)
+            {
+                var flag = extendedFlags[i];
+                var obj = spawnedObjects[i];
+                obj.AddOwnershipExtended(flag);
+                obj.SendOwnershipStatusUpdate();
+
+                m_ObjectToValidate = obj;
+                yield return WaitForConditionOrTimeOut(ValidatePermissionsOnAllClients);
+                AssertOnTimeout($"[{flag}] Not all clients have set the permission on the object!");
+            }
+
+            // Despawn the objects
+            foreach (var networkObject in spawnedObjects)
+            {
+                networkObject.Despawn(false);
+            }
+
+            yield return WaitForDespawnedOnAllOrTimeOut(spawnedObjects);
+            AssertOnTimeout("Not all clients have despawned all objects!");
+
+            // Respawn the objects
+            foreach (var networkObject in spawnedObjects)
+            {
+                networkObject.Spawn();
+            }
+
+            yield return WaitForSpawnedOnAllOrTimeOut(spawnedObjects);
+            AssertOnTimeout("[ReSpawn] Not all objects are spawned on all clients!");
+
+            // Requesting ownership should work as expected
+            var nonAuthorityIndex = 1; // Start the index at one more than the initial ownership
+            var objectToOwner = new Dictionary<NetworkObject, NetworkManager>();
+            foreach (var networkObject in spawnedObjects)
+            {
+                var nextOwner = GetNonAuthorityNetworkManager(nonAuthorityIndex++);
+                Assert.That(nextOwner.LocalClientId, Is.Not.EqualTo(networkObject.OwnerClientId));
+                Assert.IsTrue(nextOwner.SpawnManager.SpawnedObjects.TryGetValue(networkObject.NetworkObjectId, out var nextOwnerInstance));
+                Assert.That(nextOwnerInstance, Is.Not.Null);
+
+                var requestStatus = nextOwnerInstance.RequestOwnership();
+                Assert.That(requestStatus, Is.EqualTo(NetworkObject.OwnershipRequestStatus.RequestSent));
+
+                objectToOwner.Add(networkObject, nextOwner);
+            }
+
+            yield return WaitForConditionOrTimeOut(errorLog =>
+            {
+                foreach (var client in m_NetworkManagers)
+                {
+                    foreach (var (networkObject, expectedOwner) in objectToOwner)
+                    {
+                        if (!client.SpawnManager.SpawnedObjects.TryGetValue(networkObject.NetworkObjectId, out var clientInstance))
+                        {
+                            errorLog.AppendLine($"Object-{networkObject.NetworkObjectId} is not spawned on client-{client.LocalClientId}");
+                            return false;
+                        }
+
+                        if (clientInstance.OwnerClientId != expectedOwner.LocalClientId)
+                        {
+                            errorLog.AppendLine($"Client-{client.LocalClientId} has the incorrect owner for Object-{networkObject.NetworkObjectId}");
+                            return false;
+                        }
+
+                        if (clientInstance.IsOwner != (client == expectedOwner))
+                        {
+                            errorLog.AppendLine($"Client-{client.LocalClientId} has the incorrect IsOwner for Object-{networkObject.NetworkObjectId}");
+                            return false;
+                        }
+                    }
+                }
+
+                return true;
+            });
+            AssertOnTimeout("Some ownership requests failed!");
+
+            foreach (var networkObject in spawnedObjects)
+            {
+                m_ObjectToValidate = networkObject;
+                yield return WaitForConditionOrTimeOut(ValidatePermissionsOnAllClients);
+                AssertOnTimeout($"[Object-{networkObject.NetworkObjectId}] Not all clients have set the permission on the object!");
+            }
+        }
+
         internal class OwnershipPermissionsTestHelper : NetworkBehaviour
         {
             public static NetworkObject CurrentOwnedInstance;

com.unity.netcode.gameobjects/Tests/Runtime/TestHelpers/NetcodeIntegrationTest.cs

@@ -2001,6 +2001,33 @@ bool ValidateObjectsSpawnedOnAllClients(StringBuilder errorLog)
             yield return WaitForConditionOrTimeOut(ValidateObjectsSpawnedOnAllClients, timeOutHelper);
         }
 
+        /// <summary>
+        /// Waits until all given NetworkObjects are spawned on all clients or a timeout occurs.
+        /// </summary>
+        /// <param name="networkObjects">The list of <see cref="NetworkObject"/>s to wait for.</param>
+        /// <param name="timeOutHelper">An optional <see cref="TimeoutHelper"/> to control the timeout period. If null, the default timeout is used.</param>
+        /// <returns>An <see cref="IEnumerator"/> for use in Unity coroutines.</returns>
+        protected IEnumerator WaitForDespawnedOnAllOrTimeOut(List<NetworkObject> networkObjects, TimeoutHelper timeOutHelper = null)
+        {
+            bool ValidateObjectsDespawnedOnAllClients(StringBuilder errorLog)
+            {
+                foreach (var client in m_NetworkManagers)
+                {
+                    foreach (var networkObject in networkObjects)
+                    {
+                        if (client.SpawnManager.SpawnedObjects.TryGetValue(networkObject.NetworkObjectId, out NetworkObject clientObj) && clientObj.IsSpawned)
+                        {
+                            errorLog.Append($"Object-{networkObject.NetworkObjectId} is still spawned on Client-{client.LocalClientId}!");
+                            return false;
+                        }
+                    }
+                }
+                return true;
+            }
+
+            yield return WaitForConditionOrTimeOut(ValidateObjectsDespawnedOnAllClients, timeOutHelper);
+        }
+
         /// <summary>
         /// Validates that all remote clients (i.e. non-server) detect they are connected
         /// to the server and that the server reflects the appropriate number of clients