Enhance object disposal tracking and cleanup in ScopedDictionary

TUnit.Core/Data/ScopedDictionary.cs

@@ -1,4 +1,6 @@
-namespace TUnit.Core.Data;
+using TUnit.Core.Tracking;
+
+namespace TUnit.Core.Data;
 
 public class ScopedDictionary<TScope>
     where TScope : notnull
@@ -9,6 +11,24 @@ public class ScopedDictionary<TScope>
     {
         var innerDictionary = _scopedContainers.GetOrAdd(scope, _ => new GetOnlyDictionary<Type, object?>());
 
-        return innerDictionary.GetOrAdd(type, factory);
+        var obj = innerDictionary.GetOrAdd(type, factory);
+
+        ObjectTracker.OnDisposed(obj, () =>
+        {
+            innerDictionary.Remove(type);
+        });
+
+        return obj;
+    }
+
+    /// <summary>
+    /// Removes a specific value from all scopes and types where it might be stored.
+    /// This is used to clear disposed objects from the cache.
+    /// </summary>
+    public void RemoveValue(object valueToRemove)
+    {
+        // Since GetOnlyDictionary doesn't support removal, we'll need to track this differently
+        // For now, log that this needs to be implemented
+        Console.WriteLine($"[ScopedDictionary] WARNING: RemoveValue called for {valueToRemove.GetType().Name}@{valueToRemove.GetHashCode():X8} but GetOnlyDictionary doesn't support removal");
     }
 }

TUnit.Core/TestDataContainer.cs

@@ -17,13 +17,11 @@ internal static class TestDataContainer
 
     public static object? GetInstanceForAssembly(Assembly assembly, Type type, Func<Type, object> func)
     {
-        Console.WriteLine($"[TestDataContainer] GetInstanceForAssembly called for type {type.Name}, assembly {assembly.GetName().Name}");
         return _assemblyContainer.GetOrCreate(assembly, type, func);
     }
 
     public static object? GetGlobalInstance(Type type, Func<Type, object> func)
     {
-        Console.WriteLine($"[TestDataContainer] GetGlobalInstance called for type {type.Name}");
         return _globalContainer.GetOrCreate(typeof(object).FullName!, type, func);
     }
 

TUnit.Core/Tracking/ObjectTracker.cs

@@ -45,13 +45,15 @@ public static void TrackObject(TestContextEvents events, object? obj)
 
         counter.Increment();
 
+        var objType = obj.GetType().Name;
+
         events.OnDispose += async (_, _) =>
         {
             var count = counter.Decrement();
 
             if (count < 0)
             {
-                throw new InvalidOperationException($"Reference count for object {obj.GetType().Name} went below zero. This indicates a bug in the reference counting logic.");
+                throw new InvalidOperationException($"Reference count for object {objType} went below zero. This indicates a bug in the reference counting logic.");
             }
 
             if (count == 0)
@@ -68,4 +70,21 @@ private static bool ShouldSkipTracking(object? obj)
     {
         return obj is not IDisposable and not IAsyncDisposable;
     }
+
+    public static void OnDisposed(object? o, Action action)
+    {
+        if(o is not IDisposable and not IAsyncDisposable)
+        {
+            return;
+        }
+
+        _trackedObjects.GetOrAdd(o, _ => new Counter())
+            .OnCountChanged += (_, count) =>
+        {
+            if (count == 0)
+            {
+                action();
+            }
+        };
+    }
 }