Cleanups

Author: javiercn

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

@@ -398,6 +398,60 @@ public async Task RestoreStateAsync_ValidatesRegisteredUpdateCallbacksAreInvoked
         Assert.True(callbackInvoked);
     }
 
+    [Theory]
+    [MemberData(nameof(RestoreContexts))]
+    public async Task RestoreStateAsync_RestoresServicesForDifferentContexts(RestoreContext context)
+    {
+        // Arrange
+        var componentRenderMode = new TestRenderMode();
+        var (serviceProvider, persistenceManager) = CreatePersistenceManagerWithService(componentRenderMode);
+
+        var service = serviceProvider.GetRequiredService<TestPersistentService>();
+        service.TestProperty = "Test Value";
+
+        var store = new TestStore([]);
+        await persistenceManager.RestoreStateAsync(store, RestoreContext.InitialValue);
+        await persistenceManager.PersistStateAsync(store, new TestRenderer());
+
+        // Act - Create new service provider and persistence manager
+        var (newServiceProvider, newPersistenceManager) = CreatePersistenceManagerWithService(componentRenderMode);
+
+        // Restore with the specified context
+        await newPersistenceManager.RestoreStateAsync(store, context);
+
+        // Assert
+        var restoredService = newServiceProvider.GetRequiredService<TestPersistentService>();
+        Assert.Equal("Test Value", restoredService.TestProperty);
+
+        static (IServiceProvider serviceProvider, ComponentStatePersistenceManager persistenceManager) CreatePersistenceManagerWithService(IComponentRenderMode renderMode)
+        {
+            var serviceProvider = new ServiceCollection()
+                .AddScoped<TestPersistentService>()
+                .AddPersistentService<TestPersistentService>(renderMode)
+                .BuildServiceProvider();
+
+            var persistenceManager = new ComponentStatePersistenceManager(
+                NullLogger<ComponentStatePersistenceManager>.Instance,
+                serviceProvider);
+            persistenceManager.SetPlatformRenderMode(renderMode);
+
+            return (serviceProvider, persistenceManager);
+        }
+    }
+
+    public static IEnumerable<object[]> RestoreContexts()
+    {
+        yield return new object[] { RestoreContext.InitialValue };
+        yield return new object[] { RestoreContext.ValueUpdate };
+        yield return new object[] { RestoreContext.LastSnapshot };
+    }
+
+    private class TestPersistentService
+    {
+        [PersistentState(AllowUpdates = true)]
+        public string TestProperty { get; set; } = string.Empty;
+    }
+
     private class TestRenderer : Renderer
     {
         public TestRenderer() : base(new ServiceCollection().BuildServiceProvider(), NullLoggerFactory.Instance)

src/Components/Components/test/PersistentState/RestoreContextTest.cs

@@ -0,0 +1,105 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.AspNetCore.Components;
+
+public class RestoreContextTest
+{
+    [Theory]
+    [InlineData(RestoreBehavior.Default)]
+    [InlineData(RestoreBehavior.SkipLastSnapshot)]
+    public void ShouldRestore_InitialValueContext_WithDefaultOrSkipLastSnapshot(RestoreBehavior behavior)
+    {
+        var options = new RestoreOptions { RestoreBehavior = behavior };
+
+        var result = RestoreContext.InitialValue.ShouldRestore(options);
+
+        Assert.True(result);
+    }
+
+    [Fact]
+    public void ShouldRestore_InitialValueContext_WithSkipInitialValue()
+    {
+        var options = new RestoreOptions { RestoreBehavior = RestoreBehavior.SkipInitialValue };
+
+        var result = RestoreContext.InitialValue.ShouldRestore(options);
+
+        Assert.False(result);
+    }
+
+    [Theory]
+    [InlineData(RestoreBehavior.Default, true, true)]
+    [InlineData(RestoreBehavior.Default, false, true)]
+    [InlineData(RestoreBehavior.SkipInitialValue, true, false)]
+    [InlineData(RestoreBehavior.SkipInitialValue, false, false)]
+    public void ShouldRestore_InitialValueContext_ShouldRestore_IsIndependentOfAllowUpdates(RestoreBehavior behavior, bool allowUpdates, bool expectedResult)
+    {
+        var options = new RestoreOptions { RestoreBehavior = behavior, AllowUpdates = allowUpdates };
+
+        var result = RestoreContext.InitialValue.ShouldRestore(options);
+
+        Assert.Equal(expectedResult, result);
+    }
+
+    [Theory]
+    [InlineData(RestoreBehavior.Default)]
+    [InlineData(RestoreBehavior.SkipInitialValue)]
+    public void ShouldRestore_LastSnapshotContext_WithDefaultOrSkipInitialValue(RestoreBehavior behavior)
+    {
+        var options = new RestoreOptions { RestoreBehavior = behavior };
+
+        var result = RestoreContext.LastSnapshot.ShouldRestore(options);
+
+        Assert.True(result);
+    }
+
+    [Fact]
+    public void ShouldRestore_LastSnapshotContext_WithSkipLastSnapshot()
+    {
+        var options = new RestoreOptions { RestoreBehavior = RestoreBehavior.SkipLastSnapshot };
+
+        var result = RestoreContext.LastSnapshot.ShouldRestore(options);
+
+        Assert.False(result);
+    }
+
+    [Theory]
+    [InlineData(RestoreBehavior.Default, true, true)]
+    [InlineData(RestoreBehavior.Default, false, true)]
+    [InlineData(RestoreBehavior.SkipLastSnapshot, true, false)]
+    [InlineData(RestoreBehavior.SkipLastSnapshot, false, false)]
+    public void ShouldRestore_LastSnapshotContext_ShouldRestore_IsIndependentOfAllowUpdates(RestoreBehavior behavior, bool allowUpdates, bool expectedResult)
+    {
+        var options = new RestoreOptions { RestoreBehavior = behavior, AllowUpdates = allowUpdates };
+
+        var result = RestoreContext.LastSnapshot.ShouldRestore(options);
+
+        Assert.Equal(expectedResult, result);
+    }
+
+    [Theory]
+    [InlineData(RestoreBehavior.Default)]
+    [InlineData(RestoreBehavior.SkipInitialValue)]
+    [InlineData(RestoreBehavior.SkipLastSnapshot)]
+    public void ShouldRestore_ValueUpdateContext_WithoutAllowUpdates(RestoreBehavior behavior)
+    {
+        var options = new RestoreOptions { RestoreBehavior = behavior };
+
+        var result = RestoreContext.ValueUpdate.ShouldRestore(options);
+
+        Assert.False(result);
+    }
+
+    [Theory]
+    [InlineData(RestoreBehavior.Default)]
+    [InlineData(RestoreBehavior.SkipInitialValue)]
+    [InlineData(RestoreBehavior.SkipLastSnapshot)]
+    public void ShouldRestore_ValueUpdateContext_WithAllowUpdates(RestoreBehavior behavior)
+    {
+        var options = new RestoreOptions { AllowUpdates = true, RestoreBehavior = behavior };
+
+        var result = RestoreContext.ValueUpdate.ShouldRestore(options);
+
+        Assert.True(result);
+    }
+}