Make AssemblyInitialize/AssemblyCleanup inheritable

Youssef1313 · Youssef1313 · commit 6b7aa690fd79 · 2025-01-16T12:22:07.000+01:00
diff --git a/src/Adapter/MSTest.TestAdapter/Execution/TestAssemblyInfo.cs b/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,
diff --git a/src/Adapter/MSTest.TestAdapter/Execution/TypeCache.cs b/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)
             {
@@ -571,7 +595,7 @@ private void UpdateInfoIfClassInitializeOrCleanupMethod(
             }
         }
 
-        if (isCleanupMethod)
+        if (classCleanupAttribute is not null)
         {
             if (TryGetTimeoutInfo(methodInfo, FixtureKind.ClassCleanup) is { } timeoutInfo)
             {
diff --git a/src/Adapter/MSTest.TestAdapter/Extensions/MethodInfoExtensions.cs b/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")]
diff --git a/src/TestFramework/TestFramework/Attributes/Lifecycle/Cleanup/AssemblyCleanupAttribute.cs b/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();
+}
diff --git a/src/TestFramework/TestFramework/Attributes/Lifecycle/Cleanup/AssemblyCleanupExecutionContext.cs b/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; }
+}
diff --git a/src/TestFramework/TestFramework/Attributes/Lifecycle/Initialization/AssemblyInitializeAttribute.cs b/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();
+}
diff --git a/src/TestFramework/TestFramework/Attributes/Lifecycle/Initialization/AssemblyInitializeExecutionContext.cs b/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; }
+}
diff --git a/src/TestFramework/TestFramework/PublicAPI/PublicAPI.Unshipped.txt b/src/TestFramework/TestFramework/PublicAPI/PublicAPI.Unshipped.txt
@@ -1,5 +1,11 @@
 #nullable enable
 abstract Microsoft.VisualStudio.TestTools.UnitTesting.RetryBaseAttribute.ExecuteAsync(Microsoft.VisualStudio.TestTools.UnitTesting.RetryContext retryContext) -> System.Threading.Tasks.Task<Microsoft.VisualStudio.TestTools.UnitTesting.RetryResult!>!
+Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyCleanupExecutionContext
+Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyCleanupExecutionContext.AssemblyCleanupExecutionContext() -> void
+Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyCleanupExecutionContext.AssemblyCleanupExecutorGetter.get -> System.Func<System.Threading.Tasks.Task!>!
+Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyInitializeExecutionContext
+Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyInitializeExecutionContext.AssemblyInitializeExecutionContext() -> void
+Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyInitializeExecutionContext.AssemblyInitializeExecutorGetter.get -> System.Func<System.Threading.Tasks.Task!>!
 Microsoft.VisualStudio.TestTools.UnitTesting.Assert.AssertAreEqualInterpolatedStringHandler<TArgument>
 Microsoft.VisualStudio.TestTools.UnitTesting.Assert.AssertAreEqualInterpolatedStringHandler<TArgument>.AppendFormatted(object? value, int alignment = 0, string? format = null) -> void
 Microsoft.VisualStudio.TestTools.UnitTesting.Assert.AssertAreEqualInterpolatedStringHandler<TArgument>.AppendFormatted(string? value) -> void
@@ -266,3 +272,5 @@ static Microsoft.VisualStudio.TestTools.UnitTesting.Assert.ThrowsExactly<TExcept
 static Microsoft.VisualStudio.TestTools.UnitTesting.Assert.ThrowsExactly<TException>(System.Action! action, System.Func<System.Exception?, string!>! messageBuilder) -> TException!
 static Microsoft.VisualStudio.TestTools.UnitTesting.Assert.ThrowsExactlyAsync<TException>(System.Func<System.Threading.Tasks.Task!>! action, string! message = "", params object![]! messageArgs) -> System.Threading.Tasks.Task<TException!>!
 static Microsoft.VisualStudio.TestTools.UnitTesting.Assert.ThrowsExactlyAsync<TException>(System.Func<System.Threading.Tasks.Task!>! action, System.Func<System.Exception?, string!>! messageBuilder) -> System.Threading.Tasks.Task<TException!>!
+virtual Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyCleanupAttribute.ExecuteAsync(Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyCleanupExecutionContext assemblyCleanupContext) -> System.Threading.Tasks.Task!
+virtual Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyInitializeAttribute.ExecuteAsync(Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyInitializeExecutionContext assemblyInitializeContext) -> System.Threading.Tasks.Task!

Original file line number	Diff line number	Diff line change
`@@ -51,6 +51,8 @@ internal set`
`51`	`51`	`}`
`52`	`52`	`}`
`53`	`53`
	`54`	`+ internal AssemblyInitializeAttribute? AssemblyInitializeAttribute { get; set; }`
	`55`	`+`
`54`	`56`	`/// <summary>`
`55`	`57`	`/// Gets or sets the AssemblyInitializeMethod timeout.`
`56`	`58`	`/// </summary>`
`@@ -80,6 +82,8 @@ internal set`
`80`	`82`	`}`
`81`	`83`	`}`
`82`	`84`
	`85`	`+ internal AssemblyCleanupAttribute? AssemblyCleanupAttribute { get; set; }`
	`86`	`+`
`83`	`87`	`/// <summary>`
`84`	`88`	`/// Gets a value indicating whether <c>AssemblyInitialize</c> has been executed.`
`85`	`89`	`/// </summary>`
`@@ -143,7 +147,7 @@ public void RunAssemblyInitialize(TestContext testContext)`
`143`	`147`	`try`
`144`	`148`	`{`
`145`	`149`	`AssemblyInitializationException = FixtureMethodRunner.RunWithTimeoutAndCancellation(`
`146`		`- () => AssemblyInitializeMethod.InvokeAsSynchronousTask(null, testContext),`
	`150`	`+ () => AssemblyInitializeAttribute!.ExecuteAsync(new AssemblyInitializeExecutionContext(() => AssemblyInitializeMethod.InvokeAsync(null, testContext))).GetAwaiter().GetResult(),`
`147`	`151`	`testContext.CancellationTokenSource,`
`148`	`152`	`AssemblyInitializeMethodTimeoutMilliseconds,`
`149`	`153`	`AssemblyInitializeMethod,`
`@@ -216,7 +220,7 @@ public void RunAssemblyInitialize(TestContext testContext)`
`216`	`220`	`try`
`217`	`221`	`{`
`218`	`222`	`AssemblyCleanupException = FixtureMethodRunner.RunWithTimeoutAndCancellation(`
`219`		`- () => AssemblyCleanupMethod.InvokeAsSynchronousTask(null),`
	`223`	`+ () => AssemblyCleanupAttribute!.ExecuteAsync(new AssemblyCleanupExecutionContext(() => AssemblyCleanupMethod.InvokeAsync(null))).GetAwaiter().GetResult(),`
`220`	`224`	`new CancellationTokenSource(),`
`221`	`225`	`AssemblyCleanupMethodTimeoutMilliseconds,`
`222`	`226`	`AssemblyCleanupMethod,`
`@@ -276,11 +280,11 @@ internal void ExecuteAssemblyCleanup(TestContext testContext)`
`276`	`280`	`{`
`277`	`281`	`if (AssemblyCleanupMethod.GetParameters().Length == 0)`
`278`	`282`	`{`
`279`		`- AssemblyCleanupMethod.InvokeAsSynchronousTask(null);`
	`283`	`+ AssemblyCleanupAttribute!.ExecuteAsync(new AssemblyCleanupExecutionContext(() => AssemblyCleanupMethod.InvokeAsync(null))).GetAwaiter().GetResult();`
`280`	`284`	`}`
`281`	`285`	`else`
`282`	`286`	`{`
`283`		`- AssemblyCleanupMethod.InvokeAsSynchronousTask(null, testContext);`
	`287`	`+ AssemblyCleanupAttribute!.ExecuteAsync(new AssemblyCleanupExecutionContext(() => AssemblyCleanupMethod.InvokeAsync(null, testContext))).GetAwaiter().GetResult();`
`284`	`288`	`}`
`285`	`289`	`},`
`286`	`290`	`testContext.CancellationTokenSource,`
Original file line number	Diff line number	Diff line change
`@@ -416,14 +416,16 @@ private TestAssemblyInfo GetAssemblyInfo(Type type)`
`416`	`416`	`// Enumerate through all methods and identify the Assembly Init and cleanup methods.`
`417`	`417`	`foreach (MethodInfo methodInfo in PlatformServiceProvider.Instance.ReflectionOperations.GetDeclaredMethods(t))`
`418`	`418`	`{`
`419`		`- if (IsAssemblyOrClassInitializeMethod<AssemblyInitializeAttribute>(methodInfo))`
	`419`	`+ if (GetAssemblyOrClassInitializeMethod<AssemblyInitializeAttribute>(methodInfo) is { } assemblyInitializeAttribute)`
`420`	`420`	`{`
`421`	`421`	`assemblyInfo.AssemblyInitializeMethod = methodInfo;`
	`422`	`+ assemblyInfo.AssemblyInitializeAttribute = assemblyInitializeAttribute;`
`422`	`423`	`assemblyInfo.AssemblyInitializeMethodTimeoutMilliseconds = TryGetTimeoutInfo(methodInfo, FixtureKind.AssemblyInitialize);`
`423`	`424`	`}`
`424`		`- else if (IsAssemblyOrClassCleanupMethod<AssemblyCleanupAttribute>(methodInfo))`
	`425`	`+ else if (GetAssemblyOrClassCleanupMethod<AssemblyCleanupAttribute>(methodInfo) is { } assemblyCleanupAttribute)`
`425`	`426`	`{`
`426`	`427`	`assemblyInfo.AssemblyCleanupMethod = methodInfo;`
	`428`	`+ assemblyInfo.AssemblyCleanupAttribute = assemblyCleanupAttribute;`
`427`	`429`	`assemblyInfo.AssemblyCleanupMethodTimeoutMilliseconds = TryGetTimeoutInfo(methodInfo, FixtureKind.AssemblyCleanup);`
`428`	`430`	`}`
`429`	`431`	`}`
`@@ -440,17 +442,28 @@ private TestAssemblyInfo GetAssemblyInfo(Type type)`
`440`	`442`	`/// <typeparam name="TInitializeAttribute">The initialization attribute type. </typeparam>`
`441`	`443`	`/// <param name="methodInfo"> The method info. </param>`
`442`	`444`	`/// <returns> True if its an initialization method. </returns>`
`443`		`- private bool IsAssemblyOrClassInitializeMethod<TInitializeAttribute>(MethodInfo methodInfo)`
	`445`	`+ private TInitializeAttribute? GetAssemblyOrClassInitializeMethod<TInitializeAttribute>(MethodInfo methodInfo)`
`444`	`446`	`where TInitializeAttribute : Attribute`
`445`	`447`	`{`
`446`	`448`	`// TODO: this would be inconsistent with the codebase, but potential perf gain, issue: https://github.com/microsoft/testfx/issues/2999`
`447`	`449`	`// if (!methodInfo.IsStatic)`
`448`	`450`	`// {`
`449`	`451`	`// return false;`
`450`	`452`	`// }`
`451`		`- if (!_reflectionHelper.IsNonDerivedAttributeDefined<TInitializeAttribute>(methodInfo, false))`
	`453`	`+ IEnumerable<TInitializeAttribute> attributes = _reflectionHelper.GetDerivedAttributes<TInitializeAttribute>(methodInfo, inherit: false);`
	`454`	`+ using IEnumerator<TInitializeAttribute> enumerator = attributes.GetEnumerator();`
	`455`	`+ if (!enumerator.MoveNext())`
`452`	`456`	`{`
`453`		`- return false;`
	`457`	`+ // No attribute found.`
	`458`	`+ return null;`
	`459`	`+ }`
	`460`	`+`
	`461`	`+ TInitializeAttribute attribute = enumerator.Current;`
	`462`	`+ if (enumerator.MoveNext())`
	`463`	`+ {`
	`464`	`+ // More than one attribute found.`
	`465`	`+ string message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_MultipleAttributesOnTestMethod, methodInfo.DeclaringType!.FullName, methodInfo.Name);`
	`466`	`+ throw new TypeInspectionException(message);`
`454`	`467`	`}`
`455`	`468`
`456`	`469`	`if (!methodInfo.HasCorrectClassOrAssemblyInitializeSignature())`
`@@ -459,7 +472,7 @@ private bool IsAssemblyOrClassInitializeMethod<TInitializeAttribute>(MethodInfo`
`459`	`472`	`throw new TypeInspectionException(message);`
`460`	`473`	`}`
`461`	`474`
`462`		`- return true;`
	`475`	`+ return attribute;`
`463`	`476`	`}`
`464`	`477`
`465`	`478`	`/// <summary>`
`@@ -468,17 +481,28 @@ private bool IsAssemblyOrClassInitializeMethod<TInitializeAttribute>(MethodInfo`
`468`	`481`	`/// <typeparam name="TCleanupAttribute">The cleanup attribute type.</typeparam>`
`469`	`482`	`/// <param name="methodInfo"> The method info. </param>`
`470`	`483`	`/// <returns> True if its a cleanup method. </returns>`
`471`		`- private bool IsAssemblyOrClassCleanupMethod<TCleanupAttribute>(MethodInfo methodInfo)`
	`484`	`+ private TCleanupAttribute? GetAssemblyOrClassCleanupMethod<TCleanupAttribute>(MethodInfo methodInfo)`
`472`	`485`	`where TCleanupAttribute : Attribute`
`473`	`486`	`{`
`474`	`487`	`// TODO: this would be inconsistent with the codebase, but potential perf gain, issue: https://github.com/microsoft/testfx/issues/2999`
`475`	`488`	`// if (!methodInfo.IsStatic)`
`476`	`489`	`// {`
`477`	`490`	`// return false;`
`478`	`491`	`// }`
`479`		`- if (!_reflectionHelper.IsNonDerivedAttributeDefined<TCleanupAttribute>(methodInfo, false))`
	`492`	`+ IEnumerable<TCleanupAttribute> attributes = _reflectionHelper.GetDerivedAttributes<TCleanupAttribute>(methodInfo, inherit: false);`
	`493`	`+ using IEnumerator<TCleanupAttribute> enumerator = attributes.GetEnumerator();`
	`494`	`+ if (!enumerator.MoveNext())`
`480`	`495`	`{`
`481`		`- return false;`
	`496`	`+ // No attribute found.`
	`497`	`+ return null;`
	`498`	`+ }`
	`499`	`+`
	`500`	`+ TCleanupAttribute attribute = enumerator.Current;`
	`501`	`+ if (enumerator.MoveNext())`
	`502`	`+ {`
	`503`	`+ // More than one attribute found.`
	`504`	`+ string message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_MultipleAttributesOnTestMethod, methodInfo.DeclaringType!.FullName, methodInfo.Name);`
	`505`	`+ throw new TypeInspectionException(message);`
`482`	`506`	`}`
`483`	`507`
`484`	`508`	`if (!methodInfo.HasCorrectClassOrAssemblyCleanupSignature())`
`@@ -487,7 +511,7 @@ private bool IsAssemblyOrClassCleanupMethod<TCleanupAttribute>(MethodInfo method`
`487`	`511`	`throw new TypeInspectionException(message);`
`488`	`512`	`}`
`489`	`513`
`490`		`- return true;`
	`514`	`+ return attribute;`
`491`	`515`	`}`
`492`	`516`
`493`	`517`	`#endregion`
`@@ -546,10 +570,10 @@ private void UpdateInfoIfClassInitializeOrCleanupMethod(`
`546`	`570`	`bool isBase,`
`547`	`571`	`ref MethodInfo?[] initAndCleanupMethods)`
`548`	`572`	`{`
`549`		`- bool isInitializeMethod = IsAssemblyOrClassInitializeMethod<ClassInitializeAttribute>(methodInfo);`
`550`		`- bool isCleanupMethod = IsAssemblyOrClassCleanupMethod<ClassCleanupAttribute>(methodInfo);`
	`573`	`+ ClassInitializeAttribute? classInitializeAttribute = GetAssemblyOrClassInitializeMethod<ClassInitializeAttribute>(methodInfo);`
	`574`	`+ ClassCleanupAttribute? classCleanupAttribute = GetAssemblyOrClassCleanupMethod<ClassCleanupAttribute>(methodInfo);`
`551`	`575`
`552`		`- if (isInitializeMethod)`
	`576`	`+ if (classInitializeAttribute is not null)`
`553`	`577`	`{`
`554`	`578`	`if (TryGetTimeoutInfo(methodInfo, FixtureKind.ClassInitialize) is { } timeoutInfo)`
`555`	`579`	`{`
`@@ -571,7 +595,7 @@ private void UpdateInfoIfClassInitializeOrCleanupMethod(`
`571`	`595`	`}`
`572`	`596`	`}`
`573`	`597`
`574`		`- if (isCleanupMethod)`
	`598`	`+ if (classCleanupAttribute is not null)`
`575`	`599`	`{`
`576`	`600`	`if (TryGetTimeoutInfo(methodInfo, FixtureKind.ClassCleanup) is { } timeoutInfo)`
`577`	`601`	`{`