RewaTestClassAttribute.cs

[Serializable]
public class RewaTestClassAttribute : TestClassExtensionAttribute
{
    private static readonly Uri Uri = new Uri("urn:REWA.Core.Extensibility.RewaTestClassAttribute");

    public override Uri ExtensionId
    {
        get
        {
            return Uri;
        }
    }

    /// <summary>
    /// Maximum number of test method runs to perform until it passes. Default value is 2.
    /// </summary>
    public int MaxAttempts { get; set; }

    public RewaTestClassAttribute()
    {
        // Two attempts by default
        MaxAttempts = 2;
    }

    public override TestExtensionExecution GetExecution()
    {
        return new RewaTestExtensionExecution(MaxAttempts);
    }
}

internal class RewaTestExtensionExecution : TestExtensionExecution
{
    private readonly int _maxAttempts;

    public RewaTestExtensionExecution(int maxAttempts)
    {
        _maxAttempts = maxAttempts;
    }

    public override ITestMethodInvoker CreateTestMethodInvoker(TestMethodInvokerContext context)
    {
        return new RewaTestMethodInvoker(context, _maxAttempts);
    }

    public override void Initialize(TestExecution execution)
    {
        execution.BeforeClassInitialize += ExecutionOnBeforeClassInitialize;
        execution.BeforeClassCleanup += ExecutionOnBeforeClassCleanup;
    }

    private void ExecutionOnBeforeClassInitialize(object sender, BeforeClassInitializeEventArgs e)
    {
        BaseTestSuite.ClassSetup(e.TestContext);
    }

    private void ExecutionOnBeforeClassCleanup(object sender, BeforeClassCleanupEventArgs e)
    {
        BaseTestSuite.ClassClean();
    }

    public override void Dispose()
    {

    }
}

internal class RewaTestMethodInvoker : ITestMethodInvoker
{
    private readonly int _maxAttempts;
    private readonly TestMethodInvokerContext _context;

    public RewaTestMethodInvoker(TestMethodInvokerContext context, int maxAttempts)
    {
        _context = context;

        object[] maxAttemptsAttributes = context.TestMethodInfo.GetCustomAttributes(typeof(MaxAttemptsAttribute), false);
        
        _maxAttempts = maxAttemptsAttributes.Length > 0
            ? ((MaxAttemptsAttribute) maxAttemptsAttributes[0]).Value
            : maxAttempts;
        
        if (_maxAttempts < 1)
        {
            _maxAttempts = 1;
        }
    }

    public TestMethodInvokerResult Invoke(params object[] parameters)
    {
        TestMethodInvokerResult result;
        int attempt = 0;
        
        do
        {
            attempt++;
            BaseTestSuite.OnBeforeTestMethodAttempt(attempt, _context);
            result = _context.InnerInvoker.Invoke(parameters);
            BaseTestSuite.OnAfterTestMethodAttempt(result, _context.TestContext, attempt, attempt >= _maxAttempts);

            // test attempt cleanup
            var testClassType = _context.TestMethodInfo.DeclaringType;
            if (testClassType != null)
            {
                var attemptCleanupMethod = testClassType
                    .GetMethods()
                    .FirstOrDefault(m => m.GetCustomAttributes(typeof(TestAttemptCleanupAttribute), false).Length > 0);
                if (attemptCleanupMethod != null)
                {
                    attemptCleanupMethod.Invoke(null, new object[] { result.Exception == null });
                }
            }
        }
        while (result.Exception != null && attempt < _maxAttempts);
        
        _context.TestContext.WriteLine("Test method was run {0} time(s).", attempt);

        return result;
    }
}