Profiling C# Unit Tests

When writing unit tests, sometimes we want to profile the execution of each test methods to identify memory leaks or performance regressions. We can achieve this by creating custom attributes that hooks into the test execution lifecycle.

Let’s define a hook interface as such:

interface IHook
{
    void Before(ITestMethod method);
    void After(ITestMethod method, TestResult[] result);
}

Now we can create a custom TestMethodAttribute that will use this hook to perform actions before and after each test method execution:

[AttributeUsage(AttributeTargets.Method, AllowMultiple = false)]
public class HookTestMethodAttribute(Type type) : TestMethodAttribute
{
    public override TestResult[] Execute(ITestMethod method)
    {
        IHook hook;
        if (type != null && typeof(IHook).IsAssignableFrom(type))
        {
            hook = (IHook)Activator.CreateInstance(type);
        }

        hook?.Before(method);
        TestResult[] r = [method.Invoke(null)];
        hook?.After(method, r);
        return r;
    }
}

Finally, we can create a TestClassAttribute that will dispatch the appropriate hook based on some criteria, such as an environment variable:

[AttributeUsage(AttributeTargets.Class, AllowMultiple = false)]
public class HookTestClassAttribute : TestClassAttribute
{
    private static readonly Dictionary<string, Type> HOOKS = new Dictionary<string, Type>(StringComparer.OrdinalIgnoreCase)
    {
        [ "Foo" ] = typeof(FooHook),
    };

    public override TestMethodAttribute GetTestMethodAttribute(ITestMethod attribute)
    {
        var name = Environment.GetEnvironmentVariable("HOOK_TYPE");
        if (string.IsNullOrEmpty(name) || !HOOKS.TryGet(name, out var hook))
        {
            return attribute;
        }

        return new HookTestMethodAttribute(hook);
    }
}

Of course this is just one the many ways you could dispatch hooks, other options are passing the hook type directly as a parameter to the attribute constructor:

[HookTestClass<HookType>]
public class MyTests
{
}

[HookTestClass(Hook = typeof(HookType))]
public class MyTests
{
}

We can now implement our custom hook:

public class FooHook : IHook
{
    private Stopwatch _stopwatch;

    public void Before(ITestMethod method)
    {
        _stopwatch = Stopwatch.StartNew();
    }

    public void After(ITestMethod method, TestResult[] result)
    {
        _stopwatch.Stop();
        Console.WriteLine($"Test {method.Name} took {_stopwatch.ElapsedMilliseconds} ms");
    }
}