Fix TempData lazy loading bugs (#65722)

Author: dariatiurina

src/Components/Endpoints/src/TempData/TempData.cs

@@ -49,8 +49,9 @@ public object? this[string key]
 
     public object? Get(string key)
     {
+        var value = Data.GetValueOrDefault(key);
         _retainedKeys.Remove(key);
-        return Data.GetValueOrDefault(key);
+        return value;
     }
 
     public object? Peek(string key)
@@ -60,7 +61,7 @@ public object? this[string key]
 
     public void Keep()
     {
-        _retainedKeys.UnionWith(_data.Keys);
+        _retainedKeys.UnionWith(Data.Keys);
     }
 
     public void Keep(string key)
@@ -78,8 +79,9 @@ public bool ContainsKey(string key)
 
     public bool Remove(string key)
     {
+        var removed = Data.Remove(key);
         _retainedKeys.Remove(key);
-        return Data.Remove(key);
+        return removed;
     }
 
     public IDictionary<string, object?> Save()
@@ -175,7 +177,7 @@ sealed class TempDataEnumerator : IEnumerator<KeyValuePair<string, object?>>
         public TempDataEnumerator(TempData tempData)
         {
             _tempData = tempData;
-            _innerEnumerator = tempData._data.GetEnumerator();
+            _innerEnumerator = tempData.Data.GetEnumerator();
         }
 
         public KeyValuePair<string, object?> Current

src/Components/Endpoints/test/TempData/TempDataTest.cs

@@ -250,4 +250,80 @@ public void TempData_Loads_WhenLoaded()
         Assert.True(tempData.WasLoaded);
         Assert.Equal("Value", value);
     }
+
+    [Fact]
+    public void Get_ConsumesValue_WhenFirstAccessTriggersLazyLoad()
+    {
+        var tempData = new TempData(() => new Dictionary<string, object>
+        {
+            ["Message"] = "hello"
+        });
+
+        var value = tempData.Get("Message");
+        var saved = tempData.Save();
+
+        Assert.Equal("hello", value);
+        Assert.Empty(saved);
+    }
+
+    [Fact]
+    public void Indexer_ConsumesValue_WhenFirstAccessTriggersLazyLoad()
+    {
+        var tempData = new TempData(() => new Dictionary<string, object>
+        {
+            ["Message"] = "hello"
+        });
+
+        var value = tempData["Message"];
+        var saved = tempData.Save();
+
+        Assert.Equal("hello", value);
+        Assert.Empty(saved);
+    }
+
+    [Fact]
+    public void Remove_RemovesValue_WhenFirstAccessTriggersLazyLoad()
+    {
+        var tempData = new TempData(() => new Dictionary<string, object>
+        {
+            ["Message"] = "hello"
+        });
+
+        var result = tempData.Remove("Message");
+        var saved = tempData.Save();
+
+        Assert.True(result);
+        Assert.Empty(saved);
+    }
+
+    [Fact]
+    public void Keep_RetainsAllKeys_WhenFirstAccessTriggersLazyLoad()
+    {
+        var tempData = new TempData(() => new Dictionary<string, object>
+        {
+            ["Key1"] = "Value1",
+            ["Key2"] = "Value2"
+        });
+
+        tempData.Keep();
+        var saved = tempData.Save();
+
+        Assert.Equal(2, saved.Count);
+        Assert.Equal("Value1", saved["Key1"]);
+        Assert.Equal("Value2", saved["Key2"]);
+    }
+
+    [Fact]
+    public void Enumerator_IteratesAllKeys_WhenFirstAccessTriggersLazyLoad()
+    {
+        var tempData = new TempData(() => new Dictionary<string, object>
+        {
+            ["Key1"] = "Value1",
+            ["Key2"] = "Value2"
+        });
+
+        var items = ((IEnumerable<KeyValuePair<string, object>>)tempData).ToList();
+
+        Assert.Equal(2, items.Count);
+    }
 }

src/Components/test/testassets/Components.TestServer/RazorComponents/Pages/TempData/TempDataComponent.razor

@@ -65,13 +65,6 @@
 
     protected override void OnInitialized()
     {
-        _containsMessageKey = TempData?.ContainsKey("Message") ?? false;
-        _containsPeekedValueKey = TempData?.ContainsKey("PeekedValue") ?? false;
-
-        if (Handler is not null)
-        {
-            return;
-        }
         _message = TempData!.Get("Message") as string ?? "No message";
         _peekedValue = TempData!.Peek("PeekedValue") as string ?? "No peeked value";
         _keptValue = TempData!.Get("KeptValue") as string ?? "No kept value";
@@ -82,6 +75,9 @@
         _stringArrayResult = stringArray is string[] arr ? string.Join(",", arr) : $"Wrong type: {stringArray?.GetType().Name ?? "null"}";
         _intArrayResult = intArray is int[] nums ? string.Join(",", nums) : $"Wrong type: {intArray?.GetType().Name ?? "null"}";
 
+        _containsMessageKey = TempData?.ContainsKey("Message") ?? false;
+        _containsPeekedValueKey = TempData?.ContainsKey("PeekedValue") ?? false;
+
         if (ValueToKeep == "all")
         {
             TempData!.Keep();