Fix tests

Author: javiercn

src/Components/Components/src/PersistentState/PersistentServicesRegistry.cs

@@ -45,6 +45,7 @@ internal void RegisterForPersistence(PersistentComponentState state)
             return;
         }
 
+        UpdateRegistrations(state);
         var subscriptions = new List<(PersistingComponentStateSubscription, RestoringComponentStateSubscription)>(
             _registrations.Length + 1);
         for (var i = 0; i < _registrations.Length; i++)
@@ -83,7 +84,7 @@ internal void RegisterForPersistence(PersistentComponentState state)
                 }, RenderMode),
                 state.RegisterOnRestoring(() =>
                 {
-                    Restore(state);
+                    RestoreRegistrationsIfAvailable(state);
                 }, new RestoreOptions { RestoreBehavior = RestoreBehavior.SkipLastSnapshot })));
         }
 
@@ -109,14 +110,12 @@ private static void PersistInstanceState(object instance, Type type, PersistentC
         "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code",
         Justification = "Types registered for persistence are preserved in the API call to register them and typically live in assemblies that aren't trimmed.")]
     [DynamicDependency(LinkerFlags.JsonSerialized, typeof(PersistentServiceRegistration))]
-    internal void Restore(PersistentComponentState state)
+    private void UpdateRegistrations(PersistentComponentState state)
     {
         if (state.TryTakeFromJson<PersistentServiceRegistration[]>(_registryKey, out var registry) && registry != null)
         {
             _registrations = ResolveRegistrations(_registrations.Concat(registry));
         }
-
-        RestoreRegistrationsIfAvailable(state);
     }
 
     [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "Types registered for persistence are preserved in the API call to register them and typically live in assemblies that aren't trimmed.")]

src/Components/Components/src/Rendering/ComponentState.cs

@@ -362,7 +362,7 @@ internal ValueTask DisposeInBatchAsync(RenderBatchBuilder batchBuilder)
         {
             ref var currentFrame = ref frames.Array[i];
 
-            Debug.Assert(currentFrame.FrameType == RenderTreeFrameType.Component && currentFrame.Component == null, "GetComponentKey is being invoked too soon, ComponentState is not fully constructed.");
+            Debug.Assert(currentFrame.FrameType == RenderTreeFrameType.Component && currentFrame.Component != null, "GetComponentKey is being invoked too soon, ComponentState is not fully constructed.");
 
             if (currentFrame.FrameType != RenderTreeFrameType.Component ||
                 !ReferenceEquals(Component, currentFrame.Component))

src/Components/Components/test/PersistentState/ComponentStatePersistenceManagerTest.cs

@@ -120,6 +120,8 @@ public async Task PersistStateAsync_PersistsRegistry()
         persistenceManager.SetPlatformRenderMode(new TestRenderMode());
         var testStore = new TestStore([]);
 
+        await persistenceManager.RestoreStateAsync(new TestStore([]), RestoreContext.InitialValue);
+
         // Act
         await persistenceManager.PersistStateAsync(testStore, new TestRenderer());
 

src/Components/Components/test/PersistentState/PersistentServicesRegistryTest.cs

@@ -33,6 +33,8 @@ public async Task PersistStateAsync_PersistsServiceProperties()
 
         var registry = persistenceManager.ServicesRegistry;
 
+        await persistenceManager.RestoreStateAsync(new TestStore(new Dictionary<string, byte[]>()), RestoreContext.InitialValue);
+
         await persistenceManager.PersistStateAsync(testStore, new TestRenderer());
         var componentState = new PersistentComponentState(testStore.State, [], []);
 
@@ -74,6 +76,8 @@ public async Task PersistStateAsync_PersistsBaseServiceProperties()
         persistenceManagerOne.SetPlatformRenderMode(componentRenderMode);
         var testStore = new TestStore(new Dictionary<string, byte[]>());
 
+        await persistenceManagerOne.RestoreStateAsync(new TestStore(new Dictionary<string, byte[]>()), RestoreContext.InitialValue);
+
         await persistenceManagerOne.PersistStateAsync(testStore, new TestRenderer());
 
         var scopeTwo = serviceProviderTwo.CreateAsyncScope().ServiceProvider;
@@ -140,6 +144,7 @@ public async Task PersistStateAsync_DoesNotPersistNullServiceProperties()
             scope);
         persistenceManager.SetPlatformRenderMode(componentRenderMode);
         var testStore = new TestStore(new Dictionary<string, byte[]>());
+        await persistenceManager.RestoreStateAsync(new TestStore(new Dictionary<string, byte[]>()), RestoreContext.InitialValue);
 
         // Act
         await persistenceManager.PersistStateAsync(testStore, new TestRenderer());
@@ -210,14 +215,16 @@ public async Task PersistStateAsync_RestoresStateForPersistedRegistrations()
         persistenceManagerOne.SetPlatformRenderMode(componentRenderMode);
         var testStore = new TestStore(new Dictionary<string, byte[]>());
 
+        await persistenceManagerOne.RestoreStateAsync(new TestStore(new Dictionary<string, byte[]>()), RestoreContext.InitialValue);
+
         await persistenceManagerOne.PersistStateAsync(testStore, new TestRenderer());
 
         var scopeTwo = serviceProviderTwo.CreateAsyncScope().ServiceProvider;
         var persistenceManagerTwo = new ComponentStatePersistenceManager(
             NullLogger<ComponentStatePersistenceManager>.Instance,
             scopeTwo);
 
-        await persistenceManagerTwo.RestoreStateAsync(new TestStore(testStore.State));
+        await persistenceManagerTwo.RestoreStateAsync(new TestStore(testStore.State), RestoreContext.InitialValue);
 
         // Assert
         var derivedTwo = scopeTwo.GetRequiredService<BaseService>() as DerivedTwo;
@@ -324,6 +331,8 @@ public async Task PersistStateAsync_PersistsMultipleServicesWithDifferentStates(
         persistenceManager.SetPlatformRenderMode(componentRenderMode);
         var testStore = new TestStore(new Dictionary<string, byte[]>());
 
+        await persistenceManager.RestoreStateAsync(new TestStore(new Dictionary<string, byte[]>()), RestoreContext.InitialValue);
+
         await persistenceManager.PersistStateAsync(testStore, new TestRenderer());
 
         var secondScope = serviceProvider.CreateAsyncScope().ServiceProvider;
@@ -368,6 +377,7 @@ public async Task PersistStateAsync_PersistsServiceWithComplexState()
             scope);
         persistenceManager.SetPlatformRenderMode(componentRenderMode);
         var testStore = new TestStore(new Dictionary<string, byte[]>());
+        await persistenceManager.RestoreStateAsync(new TestStore(new Dictionary<string, byte[]>()), RestoreContext.InitialValue);
 
         await persistenceManager.PersistStateAsync(testStore, new TestRenderer());
 

src/Components/Components/test/PersistentStateValueProviderKeyResolverTests.cs

@@ -110,7 +110,7 @@ public async Task PersistAsync_CanPersistMultipleComponentsOfSameType_SupportsDi
             NullLogger<ComponentStatePersistenceManager>.Instance,
             new ServiceCollection().BuildServiceProvider());
 
-        persistenceManager.State.InitializeExistingState(state, RestoreContext.InitialValue);
+        await persistenceManager.RestoreStateAsync(new TestStore([]), RestoreContext.InitialValue);
 
         var renderer = new TestRenderer();
         var component1 = new TestComponent { State = "state1" };