import org.junit.jupiter.api.extension.*;
import org.junit.jupiter.api.extension.ExtensionContext.Namespace;

import java.util.logging.Logger;

import static java.lang.System.currentTimeMillis;

/**
 * Based on TimingExtension from JUnit 5 user guide, and BenchmarkExtension from Nicolai Parlog
 * https://github.com/CodeFX-org/demo-junit-5/blob/master/src/main/java/org/codefx/demo/junit5/BenchmarkExtension.java
 */
public class BenchmarkExtension implements BeforeAllCallback, BeforeTestExecutionCallback, AfterTestExecutionCallback, AfterAllCallback {
    
    private static final Namespace NAMESPACE = Namespace.create( BenchmarkExtension.class );
    private static final Logger LOGGER = Logger.getLogger( BenchmarkExtension.class.getName() );
        
    @Override
    public void beforeAll( ExtensionContext context ) {
        context.getStore( NAMESPACE ).put( context.getRequiredTestClass(), currentTimeMillis() );
    }
    
    @Override
    public void afterAll( ExtensionContext context ) {
        long launchTime  = context.getStore( NAMESPACE ).get( context.getRequiredTestClass(), long.class );
        long elapsedTime = currentTimeMillis() - launchTime;
        report( "Test class", context, elapsedTime );
    }
    
    @Override
    public void beforeTestExecution( ExtensionContext context ) {
        context.getStore( NAMESPACE ).put( context.getRequiredTestMethod(), currentTimeMillis() );
    }
    
    @Override
    public void afterTestExecution( ExtensionContext context ) {
        long launchTime  = context.getStore( NAMESPACE ).get( context.getRequiredTestMethod(), long.class );
        long elapsedTime = currentTimeMillis() - launchTime;
        report( "Test", context, elapsedTime );
    }
    
    private static void report( String unit, ExtensionContext context, long elapsedTime ) {
        String message = String.format( "%s '%s' took %d ms.", unit, context.getDisplayName(), elapsedTime );
        context.publishReportEntry( "Benchmark", message );
        LOGGER.info( message );
    }
    
}