Add support for TestInitMethod and TestCleanupMethod

Source/Test/Rewriting/Passes/MSTestRewriter.cs

@@ -19,13 +19,64 @@ internal class MSTestRewriter : AssemblyRewriter
         /// </summary>
         private readonly Configuration Configuration;
 
+        /// <summary>
+        /// MethodReference of TestInitialize method within a test class.
+        /// </summary>
+        private MethodDefinition TestInitMethod;
+        private const string TestInitAttribute = "Microsoft.VisualStudio.TestTools.UnitTesting.TestInitializeAttribute";
+
+        /// <summary>
+        /// MethodReference of TestCleanupMethod method within a test class.
+        /// </summary>
+        private MethodDefinition TestCleanupMethod;
+        private const string TestCleanupAttribute = "Microsoft.VisualStudio.TestTools.UnitTesting.TestCleanupAttribute";
+
+        /// <summary>
+        /// TestClass Attribute of MSTests.
+        /// </summary>
+        private const string TestClassAttribute = "Microsoft.VisualStudio.TestTools.UnitTesting.TestClassAttribute";
+
         /// <summary>
         /// Initializes a new instance of the <see cref="MSTestRewriter"/> class.
         /// </summary>
         internal MSTestRewriter(Configuration configuration, ILogger logger)
             : base(logger)
         {
             this.Configuration = configuration;
+            this.TestInitMethod = null;
+            this.TestCleanupMethod = null;
+        }
+
+        /// <inheritdoc/>
+        internal override void VisitType(TypeDefinition type)
+        {
+            if (type.IsAbstract)
+            {
+                return;
+            }
+
+            this.TestInitMethod = null;
+            this.TestCleanupMethod = null;
+
+            if (type.CustomAttributes.Any(p => p.AttributeType.FullName == TestClassAttribute))
+            {
+                foreach (var method in type.Methods)
+                {
+                    if (method.CustomAttributes.Count > 0)
+                    {
+                        // Assuming that we can have at max only one TestInit and TestCleanup Methods.
+                        if (method.CustomAttributes.Any(p => p.AttributeType.FullName == TestCleanupAttribute))
+                        {
+                            this.TestCleanupMethod = method;
+                        }
+
+                        if (method.CustomAttributes.Any(p => p.AttributeType.FullName == TestInitAttribute))
+                        {
+                            this.TestInitMethod = method;
+                        }
+                    }
+                }
+            }
         }
 
         /// <inheritdoc/>
@@ -170,6 +221,8 @@ internal void RewriteTestMethod(MethodDefinition method, MethodDefinition testMe
             // The emitted IL corresponds to a method body such as:
             //   Configuration configuration = Configuration.Create();
             //   TestingEngine engine = TestingEngine.Create(configuration, new Action(Test));
+            //   [engine.RegisterPerIterationInitMethod(new Action(method));]
+            //   [engine.RegisterPerIterationCallBack(new Action(method));]
             //   engine.Run();
             //   engine.ThrowIfBugFound();
             //
@@ -253,6 +306,35 @@ internal void RewriteTestMethod(MethodDefinition method, MethodDefinition testMe
             processor.Emit(OpCodes.Newobj, actionConstructor);
             processor.Emit(OpCodes.Call, createEngineMethod);
             processor.Emit(OpCodes.Dup);
+
+            // Add call to engine.RegisterPerIterationInitMethod(new Action(method));
+            if (this.TestInitMethod != null)
+            {
+                processor.Emit(OpCodes.Dup);
+                processor.Emit(OpCodes.Ldarg_0);
+                processor.Emit(OpCodes.Ldftn, this.TestInitMethod);
+                processor.Emit(OpCodes.Newobj, actionConstructor);
+
+                MethodReference registerPerIterationInitMethod = this.Module.ImportReference(
+                FindMatchingMethodInDeclaringType(resolvedEngineType, "RegisterPerIterationInitMethod", actionType));
+
+                processor.Emit(OpCodes.Call, registerPerIterationInitMethod);
+            }
+
+            // Add call to engine.RegisterPerIterationCallBack(new Action(method));
+            if (this.TestCleanupMethod != null)
+            {
+                processor.Emit(OpCodes.Dup);
+                processor.Emit(OpCodes.Ldarg_0);
+                processor.Emit(OpCodes.Ldftn, this.TestCleanupMethod);
+                processor.Emit(OpCodes.Newobj, actionConstructor);
+
+                MethodReference registerPerIterationCallBack = this.Module.ImportReference(
+                FindMatchingMethodInDeclaringType(resolvedEngineType, "RegisterPerIterationCallBack", actionType));
+
+                processor.Emit(OpCodes.Call, registerPerIterationCallBack);
+            }
+
             this.EmitMethodCall(processor, resolvedEngineType, "Run");
             this.EmitMethodCall(processor, resolvedEngineType, "ThrowIfBugFound");
             processor.Emit(OpCodes.Ret);

Source/Test/SystematicTesting/TestingEngine.cs

@@ -51,7 +51,13 @@ public sealed class TestingEngine
         /// Set of callbacks to invoke at the end
         /// of each iteration.
         /// </summary>
-        private readonly ISet<Action<uint>> PerIterationCallbacks;
+        private readonly ISet<Action> PerIterationCallbacks;
+
+        /// <summary>
+        /// Set of callbacks to invoke at the beginning
+        /// of each iteration.
+        /// </summary>
+        private readonly ISet<Action> PerIterationInitializationCallbacks;
 
         /// <summary>
         /// The scheduler used by the runtime during testing.
@@ -247,7 +253,8 @@ private TestingEngine(Configuration configuration, TestMethodInfo testMethodInfo
             this.Logger = this.DefaultLogger;
             this.Profiler = new Profiler();
 
-            this.PerIterationCallbacks = new HashSet<Action<uint>>();
+            this.PerIterationCallbacks = new HashSet<Action>();
+            this.PerIterationInitializationCallbacks = new HashSet<Action>();
 
             this.TestReport = new TestReport(configuration);
             this.ReadableTrace = string.Empty;
@@ -529,6 +536,15 @@ private bool RunNextIteration(uint iteration)
 
                 this.InitializeCustomActorLogging(runtime.DefaultActorExecutionContext);
 
+                // Invoke TestInit methods (if any) for every iteration, except the first one.
+                if (iteration != 0)
+                {
+                    foreach (var callback in this.PerIterationInitializationCallbacks)
+                    {
+                        callback();
+                    }
+                }
+
                 // Runs the test and waits for it to terminate.
                 runtime.RunTest(this.TestMethodInfo.Method, this.TestMethodInfo.Name);
                 runtime.WaitAsync().Wait();
@@ -539,7 +555,7 @@ private bool RunNextIteration(uint iteration)
                 // Invoke the per iteration callbacks, if any.
                 foreach (var callback in this.PerIterationCallbacks)
                 {
-                    callback(iteration);
+                    callback();
                 }
 
                 if (!runtime.IsBugFound)
@@ -712,14 +728,21 @@ public IEnumerable<string> TryEmitTraces(string directory, string file)
         }
 
         /// <summary>
-        /// Registers a callback to invoke at the end of each iteration. The callback takes as
-        /// a parameter an integer representing the current iteration.
+        /// Registers a callback to invoke at the end of each iteration.
         /// </summary>
-        public void RegisterPerIterationCallBack(Action<uint> callback)
+        public void RegisterPerIterationCallBack(Action callback)
         {
             this.PerIterationCallbacks.Add(callback);
         }
 
+        /// <summary>
+        /// Registers a callback to invoke at the end of each iteration.
+        /// </summary>
+        public void RegisterPerIterationInitMethod(Action callback)
+        {
+            this.PerIterationInitializationCallbacks.Add(callback);
+        }
+
         /// <summary>
         /// Take care of handling the <see cref="Configuration"/> settings for <see cref="Configuration.CustomActorRuntimeLogType"/>,
         /// <see cref="Configuration.IsDgmlGraphEnabled"/>, and <see cref="Configuration.ReportActivityCoverage"/> by setting up the