JUnit 4 TestRule to ease SWT Display Tests (Dependencies: JUnit, SWT, Guava)

DisplayHelper.java

import static com.google.common.collect.Iterables.toArray;
import static com.google.common.collect.Lists.newLinkedList;
import static java.util.Arrays.asList;

import java.util.Collection;

import org.eclipse.swt.SWT;
import org.eclipse.swt.widgets.Display;
import org.eclipse.swt.widgets.Shell;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

public class DisplayHelper implements TestRule {

  private final Collection<Shell> capturedShells;

  private boolean displayOwner;
  private Display display;

  public DisplayHelper() {
    capturedShells = newLinkedList();
    capturedShells.addAll( asList( captureShells() ) );
  }

  public Display getDisplay() {
    if( display == null ) {
      displayOwner = Display.getCurrent() == null;
      display = Display.getDefault();
    }
    return display;
  }

  public Shell[] getNewShells() {
    Collection<Shell> newShells = newLinkedList();
    Shell[] shells = captureShells();
    for( Shell shell : shells ) {
      if( !capturedShells.contains( shell ) ) {
        newShells.add( shell );
      }
    }
    return toArray( newShells, Shell.class );
  }

  public Shell createShell() {
    return createShell( SWT.NONE );
  }

  public Shell createShell( int style ) {
    return new Shell( getDisplay(), style );
  }

  public void ensureDisplay() {
    getDisplay();
  }

  public void flushPendingEvents() {
    while(    Display.getCurrent() != null
           && !Display.getCurrent().isDisposed()
           && Display.getCurrent().readAndDispatch() ) {}
  }

  public void dispose() {
    flushPendingEvents();
    disposeNewShells();
    disposeDisplay();
  }

  public Statement apply( final Statement base, Description description ) {
    return new Statement() {
      @Override
      public void evaluate() throws Throwable {
        try {
          base.evaluate();
        } finally {
          dispose();
        }
      }
    };
  }

  private void disposeNewShells() {
    Shell[] newShells = getNewShells();
    for( Shell shell : newShells ) {
      shell.dispose();
    }
  }

  private static Shell[] captureShells() {
    Shell[] result = new Shell[ 0 ];
    Display currentDisplay = Display.getCurrent();
    if( currentDisplay != null ) {
      result = currentDisplay.getShells();
    }
    return result;
  }

  private void disposeDisplay() {
    if( display != null && displayOwner ) {
      if( display.isDisposed() ) {
        display.dispose();
      }
      display = null;
    }
  }
}

DisplayHelperDemo.java

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;

import org.eclipse.swt.widgets.Display;
import org.eclipse.swt.widgets.Shell;
import org.junit.Rule;
import org.junit.Test;

public class DisplayHelperDemo {

  @Rule
  public final DisplayHelper displayHelper = new DisplayHelper();

  @Test
  public void testGetDisplay() {
    Display display = displayHelper.getDisplay();

    assertNotNull( display );
  }

  @Test
  public void testEnsureShell() {
    displayHelper.ensureDisplay();

    Display actual = Display.getCurrent();

    assertNotNull( actual );
  }

  @Test
  public void testCreateShell() {
    Shell actual = displayHelper.createShell();

    assertNotNull( actual );
  }

  @Test
  public void testGetNewShells() {
    displayHelper.createShell();

    Shell[] actuals = displayHelper.getNewShells();

    assertEquals( 1, actuals.length );
  }

  @Test
  public void testFlushPendingEvents() throws InterruptedException {
    Display display = displayHelper.getDisplay();
    int delay = 10;
    final boolean[] done = new boolean[ 1 ];
    display.timerExec( delay, new Runnable() {
      public void run() {
        done[ 0 ] = true;
      }
    } );

    Thread.sleep( delay * 2 );
    displayHelper.flushPendingEvents();

    assertTrue( done[ 0 ] );
  }

  @Test
  public void testDisplayHelperAdaptsGlobalDisplay() {
    Display expected = displayHelper.getDisplay();
    DisplayHelper localHelper = new DisplayHelper();

    Display actual = localHelper.getDisplay();


    assertSame( expected, actual );
  }

  @Test
  public void testDisplayHelperDoesNotDisposeGlobalDisplay() {
    Display display = displayHelper.getDisplay();
    DisplayHelper localHelper = new DisplayHelper();
    localHelper.ensureDisplay();

    localHelper.dispose();

    assertFalse( display.isDisposed() );
  }
}