package org.bonej.wrapperPlugins;

import net.imagej.ImageJ;
import net.imagej.ImgPlus;
import net.imagej.axis.Axes;
import net.imagej.axis.DefaultLinearAxis;
import net.imglib2.img.Img;
import net.imglib2.type.numeric.real.DoubleType;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;
import org.scijava.command.CommandModule;
import org.scijava.ui.UserInterface;
import org.scijava.ui.swing.sdi.SwingDialogPrompt;

import java.util.Iterator;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.stream.IntStream;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

/**
 * Unit tests for the {@link ConnectivityWrapper ConnectivityWrapper} plugin
 *
 * @author Richard Domander 
 */
public class ConnectivityWrapperTest {
    private static final ImageJ IMAGE_J = new ImageJ();
    private static final UserInterface mockUI = mock(UserInterface.class);
    private static final SwingDialogPrompt mockPrompt = mock(SwingDialogPrompt.class);

    @BeforeClass
    public static void oneTimeSetup() {
        IMAGE_J.ui().setDefaultUI(mockUI);
        when(mockUI.dialogPrompt(anyString(), anyString(), any(), any())).thenReturn(mockPrompt);
    }

    @AfterClass
    public static void oneTimeTearDown() {
        IMAGE_J.context().dispose();
    }

    @Test
    public void testNullImageCancelsPlugin() {
        final Future<CommandModule> future =
                IMAGE_J.command().run(ConnectivityWrapper.class, true, "inputImage", null);

        try {
            final CommandModule module = future.get();
            assertTrue("Null image should have canceled the plugin", module.isCanceled());
            assertEquals("Cancel reason is incorrect", CommonMessages.NO_IMAGE_OPEN, module.getCancelReason());
        } catch (InterruptedException | ExecutionException e) {
            e.printStackTrace();
        }
    }

    @Test
    public void test2DImageCancelsPlugin() {
        final DefaultLinearAxis xAxis = new DefaultLinearAxis(Axes.X);
        final DefaultLinearAxis yAxis = new DefaultLinearAxis(Axes.Y);
        final Img<DoubleType> img = IMAGE_J.op().create().img(new int[]{10, 10});
        final ImgPlus<DoubleType> imgPlus = new ImgPlus<>(img, "Test image", xAxis, yAxis);
        final Future<CommandModule> future =
                IMAGE_J.command().run(ConnectivityWrapper.class, true, "inputImage", imgPlus);

        try {
            final CommandModule module = future.get();
            assertTrue("2D image should have cancelled the plugin", module.isCanceled());
            assertEquals("Cancel reason is incorrect", CommonMessages.NOT_3D_IMAGE, module.getCancelReason());
        } catch (InterruptedException | ExecutionException e) {
            e.printStackTrace();
        }
    }

    @Test
    public void testNonBinaryImageCancelsPlugin() {
        // Create a test image with more than two colors
        final DefaultLinearAxis xAxis = new DefaultLinearAxis(Axes.X);
        final DefaultLinearAxis yAxis = new DefaultLinearAxis(Axes.Y);
        final DefaultLinearAxis zAxis = new DefaultLinearAxis(Axes.Z);
        final Img<DoubleType> img = IMAGE_J.op().create().img(new int[]{5, 5, 5});
        final ImgPlus<DoubleType> imgPlus = new ImgPlus<>(img, "Test image", xAxis, yAxis, zAxis);
        final Iterator<Integer> intIterator = IntStream.iterate(0, i -> i + 1).iterator();
        imgPlus.cursor().forEachRemaining(e -> e.setReal(intIterator.next()));
        final Future<CommandModule> future =
                IMAGE_J.command().run(ConnectivityWrapper.class, true, "inputImage", imgPlus);

        try {
            final CommandModule module = future.get();
            assertTrue("An image with more than two colours should have cancelled the plugin", module.isCanceled());
            assertEquals("Cancel reason is incorrect", CommonMessages.NOT_BINARY, module.getCancelReason());
        } catch (InterruptedException | ExecutionException e) {
            e.printStackTrace();
        }
    }
}
@@ -152,16 +152,149 @@
 arAxis;%0A
+import net.imagej.ops.Ops;%0Aimport net.imagej.ops.special.function.Functions;%0Aimport net.imagej.ops.special.function.UnaryFunctionOp;%0A
 import n
@@ -362,16 +362,40 @@
 leType;%0A
+import org.junit.After;%0A
 import o
@@ -411,24 +411,49 @@
 AfterClass;%0A
+import org.junit.Before;%0A
 import org.j
@@ -979,20 +979,144 @@
 Mockito.
-mock
+after;%0Aimport static org.mockito.Mockito.mock;%0Aimport static org.mockito.Mockito.reset;%0Aimport static org.mockito.Mockito.verify
 ;%0Aimport
@@ -1440,95 +1440,8 @@
 ss);
-%0A    private static final SwingDialogPrompt mockPrompt = mock(SwingDialogPrompt.class);
 %0A%0A  
@@ -1538,16 +1538,137 @@
 ockUI);%0A
+    %7D%0A%0A    @Before%0A    public void setup() %7B%0A        final SwingDialogPrompt mockPrompt = mock(SwingDialogPrompt.class);%0A
         
@@ -1760,24 +1760,94 @@
 pt);%0A    %7D%0A%0A
+    @After%0A    public void tearDown() %7B%0A        reset(mockUI);%0A    %7D%0A%0A
     @AfterCl
@@ -2396,32 +2396,125 @@
 ancelReason());%0A
+            verify(mockUI, after(100)).dialogPrompt(anyString(), anyString(), any(), any());%0A
         %7D catch 
@@ -3385,32 +3385,125 @@
 ancelReason());%0A
+            verify(mockUI, after(100)).dialogPrompt(anyString(), anyString(), any(), any());%0A
         %7D catch 
@@ -4712,32 +4712,125 @@
 ancelReason());%0A
+            verify(mockUI, after(100)).dialogPrompt(anyString(), anyString(), any(), any());%0A
         %7D catch 

Add verification