Add support for AsyncDisposable cleanup (#1288)

microsoft · Oct 6, 2022 · a2f50b2 · a2f50b2
1 parent 1626919
commit a2f50b2
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 0 deletions.
diff --git a/src/Adapter/MSTest.TestAdapter/Execution/TestMethodInfo.cs b/src/Adapter/MSTest.TestAdapter/Execution/TestMethodInfo.cs
@@ -492,6 +492,10 @@ private void RunTestCleanupMethod(object classInstance, TestResult result)
             }
             finally
             {
+#if NET6_0_OR_GREATER
+                // If you implement IAsyncDisposable without calling the DisposeAsync this would result a resource leak.
+                (classInstance as IAsyncDisposable)?.DisposeAsync().AsTask().Wait();
+#endif
                 (classInstance as IDisposable)?.Dispose();
             }
         }

diff --git a/test/UnitTests/MSTestAdapter.UnitTests/Execution/TestMethodInfoTests.cs b/test/UnitTests/MSTestAdapter.UnitTests/Execution/TestMethodInfoTests.cs
@@ -806,6 +806,46 @@ public void TestMethodInfoInvokeShouldCallDiposeForDisposableTestClass()
         Verify(disposeCalled);
     }
 
+#if NET6_0_OR_GREATER
+    public void TestMethodInfoInvoke_WhenTestClassIsAsyncDisposable_ShouldDisposeAsync()
+    {
+        //Arrange
+        var asyncDisposeCalled = false;
+        DummyTestClassWithAsyncDisposable.DisposeAsyncMethodBody = () => asyncDisposeCalled = true;
+        var ctorInfo = typeof(DummyTestClassWithAsyncDisposable).GetConstructors().Single();
+        var testClass = new TestClassInfo(typeof(DummyTestClassWithAsyncDisposable), ctorInfo, null, _classAttribute, _testAssemblyInfo);
+        var method = new TestMethodInfo(typeof(DummyTestClassWithAsyncDisposable).GetMethod("DummyTestMethod"), testClass, _testMethodOptions);
+
+        //Act
+        method.Invoke(null);
+
+        //Assert
+        Verify(asyncDisposeCalled);
+    }
+
+    public void TestMethodInfoInvoke_WhenTestClassIsDisposableAndAsyncDisposable_ShouldCallAsyncDiposeThenDipose()
+    {
+        //Arrange
+        var order = 0;
+        var disposeCalledOrder = 0;
+        var disposeAsyncCalledOrder = 0;
+
+        DummyTestClassWithAsyncDisposableAndDisposable.DisposeMethodBody = () => disposeCalledOrder = ++order;
+        DummyTestClassWithAsyncDisposableAndDisposable.DisposeAsyncMethodBody = () => disposeAsyncCalledOrder = ++order;
+
+        var ctorInfo = typeof(DummyTestClassWithAsyncDisposableAndDisposable).GetConstructors().Single();
+        var testClass = new TestClassInfo(typeof(DummyTestClassWithAsyncDisposableAndDisposable), ctorInfo, null, _classAttribute, _testAssemblyInfo);
+        var method = new TestMethodInfo(typeof(DummyTestClassWithAsyncDisposableAndDisposable).GetMethod("DummyTestMethod"), testClass, _testMethodOptions);
+
+        //Act
+        method.Invoke(null);
+
+        //Assert
+        Verify(disposeCalledOrder == 2);
+        Verify(disposeAsyncCalledOrder == 1);
+    }
+#endif
+
     public void TestMethodInfoInvokeShouldCallDiposeForDisposableTestClassIfTestCleanupThrows()
     {
         var disposeCalled = false;
@@ -1559,5 +1599,47 @@ protected override void Verify(Exception exception)
     }
 
     #endregion
+
+#if NET6_0_OR_GREATER
+    public class DummyTestClassWithAsyncDisposable : IAsyncDisposable
+    {
+        public static Action DisposeAsyncMethodBody { get; set; }
+
+        public static Action<DummyTestClassWithDisposable> DummyTestCleanupMethodBody { get; set; }
+
+        public void DummyTestMethod()
+        {
+        }
+
+        public ValueTask DisposeAsync()
+        {
+            DisposeAsyncMethodBody();
+            return ValueTask.CompletedTask;
+        }
+    }
+
+    public class DummyTestClassWithAsyncDisposableAndDisposable : IAsyncDisposable, IDisposable
+    {
+        public static Action DisposeMethodBody { get; set; }
+        public static Action DisposeAsyncMethodBody { get; set; }
+        public static Action<DummyTestClassWithDisposable> DummyTestCleanupMethodBody { get; set; }
+
+        public void DummyTestMethod()
+        {
+        }
+
+        public ValueTask DisposeAsync()
+        {
+            DisposeAsyncMethodBody();
+            return ValueTask.CompletedTask;
+        }
+
+        public void Dispose()
+        {
+            DisposeMethodBody();
+        }
+    }
+#endif
+
 }
 #endregion