Make AssemblyInitialize/AssemblyCleanup inheritable

src/Adapter/MSTest.TestAdapter/Execution/TestAssemblyInfo.cs

@@ -51,6 +51,8 @@ internal set
         }
     }
 
+    internal AssemblyInitializeAttribute? AssemblyInitializeAttribute { get; set; }
+
     /// <summary>
     /// Gets or sets the AssemblyInitializeMethod timeout.
     /// </summary>
@@ -80,6 +82,8 @@ internal set
         }
     }
 
+    internal AssemblyCleanupAttribute? AssemblyCleanupAttribute { get; set; }
+
     /// <summary>
     /// Gets a value indicating whether <c>AssemblyInitialize</c> has been executed.
     /// </summary>
@@ -143,7 +147,7 @@ public void RunAssemblyInitialize(TestContext testContext)
                     try
                     {
                         AssemblyInitializationException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
-                            () => AssemblyInitializeMethod.InvokeAsSynchronousTask(null, testContext),
+                            () => AssemblyInitializeAttribute!.ExecuteAsync(new AssemblyInitializeExecutionContext(() => AssemblyInitializeMethod.InvokeAsync(null, testContext))).GetAwaiter().GetResult(),
                             testContext.CancellationTokenSource,
                             AssemblyInitializeMethodTimeoutMilliseconds,
                             AssemblyInitializeMethod,
@@ -216,7 +220,7 @@ public void RunAssemblyInitialize(TestContext testContext)
             try
             {
                 AssemblyCleanupException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
-                     () => AssemblyCleanupMethod.InvokeAsSynchronousTask(null),
+                     () => AssemblyCleanupAttribute!.ExecuteAsync(new AssemblyCleanupExecutionContext(() => AssemblyCleanupMethod.InvokeAsync(null))).GetAwaiter().GetResult(),
                      new CancellationTokenSource(),
                      AssemblyCleanupMethodTimeoutMilliseconds,
                      AssemblyCleanupMethod,
@@ -276,11 +280,11 @@ internal void ExecuteAssemblyCleanup(TestContext testContext)
                      {
                          if (AssemblyCleanupMethod.GetParameters().Length == 0)
                          {
-                             AssemblyCleanupMethod.InvokeAsSynchronousTask(null);
+                             AssemblyCleanupAttribute!.ExecuteAsync(new AssemblyCleanupExecutionContext(() => AssemblyCleanupMethod.InvokeAsync(null))).GetAwaiter().GetResult();
                          }
                          else
                          {
-                             AssemblyCleanupMethod.InvokeAsSynchronousTask(null, testContext);
+                             AssemblyCleanupAttribute!.ExecuteAsync(new AssemblyCleanupExecutionContext(() => AssemblyCleanupMethod.InvokeAsync(null, testContext))).GetAwaiter().GetResult();
                          }
                      },
                      testContext.CancellationTokenSource,

src/Adapter/MSTest.TestAdapter/Execution/TypeCache.cs

@@ -416,14 +416,16 @@ private TestAssemblyInfo GetAssemblyInfo(Type type)
             // Enumerate through all methods and identify the Assembly Init and cleanup methods.
             foreach (MethodInfo methodInfo in PlatformServiceProvider.Instance.ReflectionOperations.GetDeclaredMethods(t))
             {
-                if (IsAssemblyOrClassInitializeMethod<AssemblyInitializeAttribute>(methodInfo))
+                if (GetAssemblyOrClassInitializeMethod<AssemblyInitializeAttribute>(methodInfo) is { } assemblyInitializeAttribute)
                 {
                     assemblyInfo.AssemblyInitializeMethod = methodInfo;
+                    assemblyInfo.AssemblyInitializeAttribute = assemblyInitializeAttribute;
                     assemblyInfo.AssemblyInitializeMethodTimeoutMilliseconds = TryGetTimeoutInfo(methodInfo, FixtureKind.AssemblyInitialize);
                 }
-                else if (IsAssemblyOrClassCleanupMethod<AssemblyCleanupAttribute>(methodInfo))
+                else if (GetAssemblyOrClassCleanupMethod<AssemblyCleanupAttribute>(methodInfo) is { } assemblyCleanupAttribute)
                 {
                     assemblyInfo.AssemblyCleanupMethod = methodInfo;
+                    assemblyInfo.AssemblyCleanupAttribute = assemblyCleanupAttribute;
                     assemblyInfo.AssemblyCleanupMethodTimeoutMilliseconds = TryGetTimeoutInfo(methodInfo, FixtureKind.AssemblyCleanup);
                 }
             }
@@ -440,17 +442,28 @@ private TestAssemblyInfo GetAssemblyInfo(Type type)
     /// <typeparam name="TInitializeAttribute">The initialization attribute type. </typeparam>
     /// <param name="methodInfo"> The method info. </param>
     /// <returns> True if its an initialization method. </returns>
-    private bool IsAssemblyOrClassInitializeMethod<TInitializeAttribute>(MethodInfo methodInfo)
+    private TInitializeAttribute? GetAssemblyOrClassInitializeMethod<TInitializeAttribute>(MethodInfo methodInfo)
         where TInitializeAttribute : Attribute
     {
         // TODO: this would be inconsistent with the codebase, but potential perf gain, issue: https://github.com/microsoft/testfx/issues/2999
         // if (!methodInfo.IsStatic)
         // {
         //    return false;
         // }
-        if (!_reflectionHelper.IsNonDerivedAttributeDefined<TInitializeAttribute>(methodInfo, false))
+        IEnumerable<TInitializeAttribute> attributes = _reflectionHelper.GetDerivedAttributes<TInitializeAttribute>(methodInfo, inherit: false);
+        using IEnumerator<TInitializeAttribute> enumerator = attributes.GetEnumerator();
+        if (!enumerator.MoveNext())
         {
-            return false;
+            // No attribute found.
+            return null;
+        }
+
+        TInitializeAttribute attribute = enumerator.Current;
+        if (enumerator.MoveNext())
+        {
+            // More than one attribute found.
+            string message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_MultipleAttributesOnTestMethod, methodInfo.DeclaringType!.FullName, methodInfo.Name);
+            throw new TypeInspectionException(message);
         }
 
         if (!methodInfo.HasCorrectClassOrAssemblyInitializeSignature())
@@ -459,7 +472,7 @@ private bool IsAssemblyOrClassInitializeMethod<TInitializeAttribute>(MethodInfo
             throw new TypeInspectionException(message);
         }
 
-        return true;
+        return attribute;
     }
 
     /// <summary>
@@ -468,17 +481,28 @@ private bool IsAssemblyOrClassInitializeMethod<TInitializeAttribute>(MethodInfo
     /// <typeparam name="TCleanupAttribute">The cleanup attribute type.</typeparam>
     /// <param name="methodInfo"> The method info. </param>
     /// <returns> True if its a cleanup method. </returns>
-    private bool IsAssemblyOrClassCleanupMethod<TCleanupAttribute>(MethodInfo methodInfo)
+    private TCleanupAttribute? GetAssemblyOrClassCleanupMethod<TCleanupAttribute>(MethodInfo methodInfo)
         where TCleanupAttribute : Attribute
     {
         // TODO: this would be inconsistent with the codebase, but potential perf gain, issue: https://github.com/microsoft/testfx/issues/2999
         // if (!methodInfo.IsStatic)
         // {
         //    return false;
         // }
-        if (!_reflectionHelper.IsNonDerivedAttributeDefined<TCleanupAttribute>(methodInfo, false))
+        IEnumerable<TCleanupAttribute> attributes = _reflectionHelper.GetDerivedAttributes<TCleanupAttribute>(methodInfo, inherit: false);
+        using IEnumerator<TCleanupAttribute> enumerator = attributes.GetEnumerator();
+        if (!enumerator.MoveNext())
         {
-            return false;
+            // No attribute found.
+            return null;
+        }
+
+        TCleanupAttribute attribute = enumerator.Current;
+        if (enumerator.MoveNext())
+        {
+            // More than one attribute found.
+            string message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_MultipleAttributesOnTestMethod, methodInfo.DeclaringType!.FullName, methodInfo.Name);
+            throw new TypeInspectionException(message);
         }
 
         if (!methodInfo.HasCorrectClassOrAssemblyCleanupSignature())
@@ -487,7 +511,7 @@ private bool IsAssemblyOrClassCleanupMethod<TCleanupAttribute>(MethodInfo method
             throw new TypeInspectionException(message);
         }
 
-        return true;
+        return attribute;
     }
 
     #endregion
@@ -546,10 +570,10 @@ private void UpdateInfoIfClassInitializeOrCleanupMethod(
         bool isBase,
         ref MethodInfo?[] initAndCleanupMethods)
     {
-        bool isInitializeMethod = IsAssemblyOrClassInitializeMethod<ClassInitializeAttribute>(methodInfo);
-        bool isCleanupMethod = IsAssemblyOrClassCleanupMethod<ClassCleanupAttribute>(methodInfo);
+        ClassInitializeAttribute? classInitializeAttribute = GetAssemblyOrClassInitializeMethod<ClassInitializeAttribute>(methodInfo);
+        ClassCleanupAttribute? classCleanupAttribute = GetAssemblyOrClassCleanupMethod<ClassCleanupAttribute>(methodInfo);
 
-        if (isInitializeMethod)
+        if (classInitializeAttribute is not null)
         {
             if (TryGetTimeoutInfo(methodInfo, FixtureKind.ClassInitialize) is { } timeoutInfo)
             {
@@ -558,8 +582,7 @@ private void UpdateInfoIfClassInitializeOrCleanupMethod(
 
             if (isBase)
             {
-                if (_reflectionHelper.GetFirstDerivedAttributeOrDefault<ClassInitializeAttribute>(methodInfo, inherit: true)?
-                        .InheritanceBehavior == InheritanceBehavior.BeforeEachDerivedClass)
+                if (classInitializeAttribute.InheritanceBehavior == InheritanceBehavior.BeforeEachDerivedClass)
                 {
                     initAndCleanupMethods[0] = methodInfo;
                 }
@@ -571,7 +594,7 @@ private void UpdateInfoIfClassInitializeOrCleanupMethod(
             }
         }
 
-        if (isCleanupMethod)
+        if (classCleanupAttribute is not null)
         {
             if (TryGetTimeoutInfo(methodInfo, FixtureKind.ClassCleanup) is { } timeoutInfo)
             {
@@ -580,8 +603,7 @@ private void UpdateInfoIfClassInitializeOrCleanupMethod(
 
             if (isBase)
             {
-                if (_reflectionHelper.GetFirstDerivedAttributeOrDefault<ClassCleanupAttribute>(methodInfo, inherit: true)?
-                        .InheritanceBehavior == InheritanceBehavior.BeforeEachDerivedClass)
+                if (classCleanupAttribute.InheritanceBehavior == InheritanceBehavior.BeforeEachDerivedClass)
                 {
                     initAndCleanupMethods[1] = methodInfo;
                 }

src/Adapter/MSTest.TestAdapter/Extensions/MethodInfoExtensions.cs

@@ -107,7 +107,7 @@ internal static bool IsValidReturnType(this MethodInfo method)
     }
 
     /// <summary>
-    /// Invoke a <see cref="MethodInfo"/> as a synchronous <see cref="Task"/>.
+    /// Invoke a <see cref="MethodInfo"/> as an asynchronous <see cref="Task"/>.
     /// </summary>
     /// <param name="methodInfo">
     /// <see cref="MethodInfo"/> instance.
@@ -118,7 +118,7 @@ internal static bool IsValidReturnType(this MethodInfo method)
     /// <param name="arguments">
     /// Arguments for the methodInfo invoke.
     /// </param>
-    internal static void InvokeAsSynchronousTask(this MethodInfo methodInfo, object? classInstance, params object?[]? arguments)
+    internal static async Task InvokeAsync(this MethodInfo methodInfo, object? classInstance, params object?[]? arguments)
     {
         ParameterInfo[]? methodParameters = methodInfo.GetParameters();
 
@@ -189,14 +189,29 @@ internal static void InvokeAsSynchronousTask(this MethodInfo methodInfo, object?
         // If methodInfo is an async method, wait for returned task
         if (invokeResult is Task task)
         {
-            task.GetAwaiter().GetResult();
+            await task;
         }
         else if (invokeResult is ValueTask valueTask)
         {
-            valueTask.GetAwaiter().GetResult();
+            await valueTask;
         }
     }
 
+    /// <summary>
+    /// Invoke a <see cref="MethodInfo"/> as a synchronous <see cref="Task"/>.
+    /// </summary>
+    /// <param name="methodInfo">
+    /// <see cref="MethodInfo"/> instance.
+    /// </param>
+    /// <param name="classInstance">
+    /// Instance of the on which methodInfo is invoked.
+    /// </param>
+    /// <param name="arguments">
+    /// Arguments for the methodInfo invoke.
+    /// </param>
+    internal static void InvokeAsSynchronousTask(this MethodInfo methodInfo, object? classInstance, params object?[]? arguments)
+        => InvokeAsync(methodInfo, classInstance, arguments).GetAwaiter().GetResult();
+
     // Scenarios to test:
     //
     // [DataRow(null, "Hello")]

src/TestFramework/TestFramework/Attributes/Lifecycle/Cleanup/AssemblyCleanupAttribute.cs

@@ -7,4 +7,13 @@ namespace Microsoft.VisualStudio.TestTools.UnitTesting;
 /// The assembly cleanup attribute.
 /// </summary>
 [AttributeUsage(AttributeTargets.Method)]
-public sealed class AssemblyCleanupAttribute : Attribute;
+public class AssemblyCleanupAttribute : Attribute
+{
+    /// <summary>
+    /// Executes the assembly cleanup method. Custom <see cref="AssemblyCleanupAttribute"/> implementations may
+    /// override this method to plug in custom logic for executing assembly cleanup.
+    /// </summary>
+    /// <param name="assemblyCleanupContext">A struct to hold information for executing the assembly cleanup.</param>
+    public virtual async Task ExecuteAsync(AssemblyCleanupExecutionContext assemblyCleanupContext)
+        => await assemblyCleanupContext.AssemblyCleanupExecutorGetter();
+}

src/TestFramework/TestFramework/Attributes/Lifecycle/Cleanup/AssemblyCleanupExecutionContext.cs

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.VisualStudio.TestTools.UnitTesting;
+
+/// <summary>
+/// Provides the information needed for executing assembly cleanup.
+/// This type is passed as a parameter to <see cref="AssemblyCleanupAttribute.ExecuteAsync(AssemblyCleanupExecutionContext)"/>.
+/// </summary>
+public readonly struct AssemblyCleanupExecutionContext
+{
+    internal AssemblyCleanupExecutionContext(Func<Task> assemblyCleanupExecutorGetter)
+        => AssemblyCleanupExecutorGetter = assemblyCleanupExecutorGetter;
+
+    /// <summary>
+    /// Gets the <see cref="Func{Task}"/> that returns the <see cref="Task"/> that executes the AssemblyCleanup method.
+    /// </summary>
+    public Func<Task> AssemblyCleanupExecutorGetter { get; }
+}

src/TestFramework/TestFramework/Attributes/Lifecycle/Initialization/AssemblyInitializeAttribute.cs

@@ -7,4 +7,13 @@ namespace Microsoft.VisualStudio.TestTools.UnitTesting;
 /// The assembly initialize attribute.
 /// </summary>
 [AttributeUsage(AttributeTargets.Method)]
-public sealed class AssemblyInitializeAttribute : Attribute;
+public class AssemblyInitializeAttribute : Attribute
+{
+    /// <summary>
+    /// Executes the assembly initialize method. Custom <see cref="AssemblyInitializeAttribute"/> implementations may
+    /// override this method to plug in custom logic for executing assembly initialize.
+    /// </summary>
+    /// <param name="assemblyInitializeContext">A struct to hold information for executing the assembly initialize.</param>
+    public virtual async Task ExecuteAsync(AssemblyInitializeExecutionContext assemblyInitializeContext)
+        => await assemblyInitializeContext.AssemblyInitializeExecutorGetter();
+}

src/TestFramework/TestFramework/Attributes/Lifecycle/Initialization/AssemblyInitializeExecutionContext.cs

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.VisualStudio.TestTools.UnitTesting;
+
+/// <summary>
+/// Provides the information needed for executing assembly initialize.
+/// This type is passed as a parameter to <see cref="AssemblyInitializeAttribute.ExecuteAsync(AssemblyInitializeExecutionContext)"/>.
+/// </summary>
+public readonly struct AssemblyInitializeExecutionContext
+{
+    internal AssemblyInitializeExecutionContext(Func<Task> assemblyInitializeExecutorGetter)
+        => AssemblyInitializeExecutorGetter = assemblyInitializeExecutorGetter;
+
+    /// <summary>
+    /// Gets the <see cref="Func{Task}"/> that returns the <see cref="Task"/> that executes the AssemblyInitialize method.
+    /// </summary>
+    public Func<Task> AssemblyInitializeExecutorGetter { get; }
+}

src/TestFramework/TestFramework/Attributes/Lifecycle/Initialization/ClassInitializeAttribute.cs

@@ -7,7 +7,7 @@ namespace Microsoft.VisualStudio.TestTools.UnitTesting;
 /// The class initialize attribute.
 /// </summary>
 [AttributeUsage(AttributeTargets.Method)]
-public sealed class ClassInitializeAttribute : Attribute
+public class ClassInitializeAttribute : Attribute
 {
     /// <summary>
     /// Initializes a new instance of the <see cref="ClassInitializeAttribute"/> class.
@@ -28,4 +28,12 @@ public sealed class ClassInitializeAttribute : Attribute
     /// Gets the Inheritance Behavior.
     /// </summary>
     public InheritanceBehavior InheritanceBehavior { get; }
+
+    /// <summary>
+    /// Executes the class initialize method. Custom <see cref="ClassInitializeAttribute"/> implementations may
+    /// override this method to plug in custom logic for executing class initialize.
+    /// </summary>
+    /// <param name="classInitializeContext">A struct to hold information for executing the class initialize.</param>
+    public virtual async Task ExecuteAsync(ClassInitializeExecutionContext classInitializeContext)
+        => await classInitializeContext.ClassInitializeExecutorGetter();
 }

src/TestFramework/TestFramework/Attributes/Lifecycle/Initialization/ClassInitializeExecutionContext.cs

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.VisualStudio.TestTools.UnitTesting;
+
+/// <summary>
+/// Provides the information needed for executing class initialize.
+/// This type is passed as a parameter to <see cref="ClassInitializeAttribute.ExecuteAsync(ClassInitializeExecutionContext)"/>.
+/// </summary>
+public readonly struct ClassInitializeExecutionContext
+{
+    internal ClassInitializeExecutionContext(Func<Task> classInitializeExecutorGetter)
+        => ClassInitializeExecutorGetter = classInitializeExecutorGetter;
+
+    /// <summary>
+    /// Gets the <see cref="Func{Task}"/> that returns the <see cref="Task"/> that executes the ClassInitialize method.
+    /// </summary>
+    public Func<Task> ClassInitializeExecutorGetter { get; }
+}