package arez;

import java.io.File;
import javax.annotation.Nonnull;
import org.realityforge.braincheck.GuardMessageCollector;
import org.testng.ITestContext;
import org.testng.ITestListener;
import org.testng.ITestResult;
import org.testng.TestListenerAdapter;
import static org.testng.Assert.*;

public final class MessageCollector
  extends TestListenerAdapter
  implements ITestListener
{
  @Nonnull
  private final GuardMessageCollector _messages = createCollector();

  @Override
  public void onTestStart( @Nonnull final ITestResult result )
  {
    if ( shouldCheckDiagnosticMessages() )
    {
      _messages.onTestStart();
    }
  }

  @Override
  public void onTestSuccess( @Nonnull final ITestResult result )
  {
    if ( shouldCheckDiagnosticMessages() )
    {
      _messages.onTestComplete();
    }
  }

  @Override
  public void onStart( @Nonnull final ITestContext context )
  {
    if ( shouldCheckDiagnosticMessages() )
    {
      _messages.onTestSuiteStart();
    }
  }

  @Override
  public void onFinish( @Nonnull final ITestContext context )
  {
    if ( 0 == context.getFailedTests().size() && shouldCheckDiagnosticMessages() )
    {
      _messages.onTestSuiteComplete();
    }
  }

  private boolean shouldCheckDiagnosticMessages()
  {
    return System.getProperty( "arez.check_diagnostic_messages", "true" ).equals( "true" );
  }

  @Nonnull
  private GuardMessageCollector createCollector()
  {
    final boolean saveIfChanged = "true".equals( System.getProperty( "arez.output_fixture_data", "false" ) );
    final String fixtureDir = System.getProperty( "arez.diagnostic_messages_file" );
    assertNotNull( fixtureDir,
                   "Expected System.getProperty( \"arez.diagnostic_messages_file\" ) to return location of diagnostic messages file" );
    return new GuardMessageCollector( "Arez", new File( fixtureDir ), saveIfChanged );
  }
}

@@ -97,21 +97,30 @@
 incheck.
-Guard
+AbstractTestNG
 MessageC
@@ -144,1088 +144,193 @@
 org.
-testng.ITestContext;%0Aimport org.testng.ITestListener;%0Aimport org.testng.ITestResult;%0Aimport org.testng.TestListenerAdapter;%0Aimport static org.testng.Assert.*;%0A%0Apublic final class MessageCollector%0A  extends TestListenerAdapter%0A  implements ITestListener%0A%7B%0A  @Nonnull%0A  private final GuardMessageCollector _messages = createCollector();%0A%0A  @Override%0A  public void onTestStart( @Nonnull final ITestResult result )%0A  %7B%0A    if ( shouldCheckDiagnosticMessages() )%0A    %7B%0A      _messages.onTestStart();%0A    %7D%0A  %7D%0A%0A  @Override%0A  public void onTestSuccess( @Nonnull final ITestResult result )%0A  %7B%0A    if ( shouldCheckDiagnosticMessages() )%0A    %7B%0A      _messages.onTestComplete();%0A    %7D%0A  %7D%0A%0A  @Override%0A  public void onStart( @Nonnull final ITestContext context )%0A  %7B%0A    if ( shouldCheckDiagnosticMessages() )%0A    %7B%0A      _messages.onTestSuiteStart();%0A    %7D%0A  %7D%0A%0A  @Override%0A  public void onFinish( @Nonnull final ITestContext context )%0A  %7B%0A    if ( 0 == context.getFailedTests().size() && shouldCheckDiagnosticMessages() )%0A    %7B%0A      _messages.onTestSuiteComplete();%0A    %7D%0A  %7D%0A%0A  private
+realityforge.braincheck.GuardMessageCollector;%0Aimport static org.testng.Assert.*;%0A%0Apublic final class MessageCollector%0A  extends AbstractTestNGMessageCollector%0A%7B%0A  @Override%0A  protected
  boo
@@ -484,15 +484,29 @@
 l%0A  
-private
+@Override%0A  protected
  Gua

Move to support provided by latest Braincheck