Reimplement the LaunchOrStop editor and renderer

VirtualDevices now keep track of their launch or stop progress with a
VirtualDevice.State enum property. Launch or stop the device by
indirectly calling TableModel::setValueAt with the editor and
corresponding enum value. The model now does the work of launching or
stopping the device.

Bug: 237442318
Test: LaunchOrStopButtonTableCellEditorTest
Change-Id: I1945951b21a4e44a32a54396e4bf0a1559393060

7 files changed, 62 insertions, 363 deletions

diff --git a/device-manager/src/com/android/tools/idea/devicemanager/IconButtonTableCellEditor.java b/device-manager/src/com/android/tools/idea/devicemanager/IconButtonTableCellEditor.java
index 0c5eca52d0c..d7b70d6972c 100644
--- a/device-manager/src/com/android/tools/idea/devicemanager/IconButtonTableCellEditor.java
+++ b/device-manager/src/com/android/tools/idea/devicemanager/IconButtonTableCellEditor.java
@@ -17,7 +17,6 @@ package com.android.tools.idea.devicemanager;
 
 import com.google.common.annotations.VisibleForTesting;
 import java.awt.Component;
-import javax.swing.AbstractButton;
 import javax.swing.AbstractCellEditor;
 import javax.swing.Icon;
 import javax.swing.JTable;
@@ -28,14 +27,10 @@ import org.jetbrains.annotations.Nullable;
 
 public class IconButtonTableCellEditor extends AbstractCellEditor implements TableCellEditor {
   protected final @NotNull IconButton myButton;
-  private final @Nullable Object myValue;
+  protected @Nullable Object myValue;
 
   protected IconButtonTableCellEditor() {
-    this(null);
-  }
-
-  protected IconButtonTableCellEditor(@Nullable Object value) {
-    this(value, null);
+    this(null, null);
   }
 
   protected IconButtonTableCellEditor(@Nullable Object value, @Nullable Icon icon) {
@@ -52,11 +47,6 @@ public class IconButtonTableCellEditor extends AbstractCellEditor implements Tab
   }
 
   @VisibleForTesting
-  public final @NotNull AbstractButton getButton() {
-    return myButton;
-  }
-
-  @VisibleForTesting
   public final @Nullable ChangeEvent getChangeEvent() {
     return changeEvent;
   }
diff --git a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellEditor.java b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellEditor.java
index ced1ea7cea8..eeebed9997b 100644
--- a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellEditor.java
+++ b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellEditor.java
@@ -15,159 +15,33 @@
  */
 package com.android.tools.idea.devicemanager.virtualtab;
 
-import com.android.annotations.concurrency.UiThread;
-import com.android.annotations.concurrency.WorkerThread;
-import com.android.ddmlib.EmulatorConsole;
-import com.android.ddmlib.IDevice;
-import com.android.sdklib.internal.avd.AvdInfo.AvdStatus;
-import com.android.tools.idea.avdmanager.AvdManagerConnection;
-import com.android.tools.idea.devicemanager.DeviceManagerAndroidDebugBridge;
-import com.android.tools.idea.devicemanager.DeviceManagerUsageTracker;
 import com.android.tools.idea.devicemanager.IconButtonTableCellEditor;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.util.concurrent.FluentFuture;
-import com.google.common.util.concurrent.FutureCallback;
-import com.google.common.util.concurrent.Futures;
-import com.google.common.util.concurrent.ListenableFuture;
-import com.google.wireless.android.sdk.stats.DeviceManagerEvent;
-import com.google.wireless.android.sdk.stats.DeviceManagerEvent.EventKind;
-import com.intellij.openapi.project.Project;
-import com.intellij.util.concurrency.AppExecutorUtil;
-import com.intellij.util.concurrency.EdtExecutorService;
-import icons.StudioIcons;
 import java.awt.Component;
-import java.util.function.BiConsumer;
-import java.util.function.Function;
-import java.util.function.Supplier;
 import javax.swing.JTable;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 
-// TODO The EDT, application pool threads, and who knows what other ones are accessing the AvdInfo. I don't like that.
 final class LaunchOrStopButtonTableCellEditor extends IconButtonTableCellEditor {
-  private final @Nullable Project myProject;
-  private final @NotNull DeviceManagerAndroidDebugBridge myBridge;
-  private final @NotNull NewSetEnabled myNewSetEnabled;
-  private final @NotNull Function<@NotNull IDevice, @NotNull EmulatorConsole> myGetConsole;
-  private final @NotNull Supplier<@NotNull AvdManagerConnection> myGetDefaultAvdManagerConnection;
-  private final @NotNull BiConsumer<@NotNull Throwable, @Nullable Project> myShowErrorDialog;
-
   private VirtualDevice myDevice;
 
-  @VisibleForTesting
-  interface NewSetEnabled {
-    @NotNull FutureCallback<@NotNull Object> apply(@NotNull LaunchOrStopButtonTableCellEditor editor);
-  }
-
-  @UiThread
-  LaunchOrStopButtonTableCellEditor(@Nullable Project project) {
-    this(project,
-         new DeviceManagerAndroidDebugBridge(),
-         SetEnabled::new,
-         EmulatorConsole::getConsole,
-         AvdManagerConnection::getDefaultAvdManagerConnection,
-         VirtualTabMessages::showErrorDialog);
-  }
-
-  @UiThread
-  @VisibleForTesting
-  LaunchOrStopButtonTableCellEditor(@Nullable Project project,
-                                    @NotNull DeviceManagerAndroidDebugBridge bridge,
-                                    @NotNull NewSetEnabled newSetEnabled,
-                                    @NotNull Function<@NotNull IDevice, @NotNull EmulatorConsole> getConsole,
-                                    @NotNull Supplier<@NotNull AvdManagerConnection> getDefaultAvdManagerConnection,
-                                    @NotNull BiConsumer<@NotNull Throwable, @Nullable Project> showErrorDialog) {
-    super(VirtualDevice.State.STOPPED);
-
-    myProject = project;
-    myBridge = bridge;
-    myNewSetEnabled = newSetEnabled;
-    myGetConsole = getConsole;
-    myGetDefaultAvdManagerConnection = getDefaultAvdManagerConnection;
-    myShowErrorDialog = showErrorDialog;
-
-    myButton.addActionListener(actionEvent -> {
+  LaunchOrStopButtonTableCellEditor() {
+    myButton.addActionListener(event -> {
       if (myDevice.isOnline()) {
-        stop();
+        myValue = VirtualDevice.State.STOPPING;
       }
       else {
-        launch(project);
+        myValue = VirtualDevice.State.LAUNCHING;
       }
-    });
-  }
-
-  @UiThread
-  private void stop() {
-    DeviceManagerEvent event = DeviceManagerEvent.newBuilder()
-      .setKind(EventKind.VIRTUAL_STOP_ACTION)
-      .build();
-
-    DeviceManagerUsageTracker.log(event);
-    myButton.setEnabled(false);
 
-    // noinspection UnstableApiUsage
-    FluentFuture.from(myBridge.findDevice(myProject, myDevice.getKey()))
-      .transform(this::stop, AppExecutorUtil.getAppExecutorService())
-      .addCallback(myNewSetEnabled.apply(this), EdtExecutorService.getInstance());
-  }
-
-  /**
-   * Called by an application pool thread
-   */
-  @WorkerThread
-  @SuppressWarnings("SameReturnValue")
-  private @Nullable Void stop(@Nullable IDevice device) {
-    if (device == null) {
-      throw new ErrorDialogException("Unable to stop " + myDevice,
-                                     "An error occurred stopping " + myDevice + ". To stop the device, try manually closing the " +
-                                     myDevice + " emulator window.");
-    }
-
-    myGetConsole.apply(device).kill();
-    return null;
-  }
-
-  @UiThread
-  private void launch(@Nullable Project project) {
-    DeviceManagerEvent event = DeviceManagerEvent.newBuilder()
-      .setKind(EventKind.VIRTUAL_LAUNCH_ACTION)
-      .build();
-
-    DeviceManagerUsageTracker.log(event);
-    myButton.setEnabled(false);
-
-    ListenableFuture<IDevice> future = myGetDefaultAvdManagerConnection.get().startAvd(project, myDevice.getAvdInfo());
-    Futures.addCallback(future, myNewSetEnabled.apply(this), EdtExecutorService.getInstance());
+      fireEditingStopped();
+    });
   }
 
   @VisibleForTesting
-  static final class SetEnabled implements FutureCallback<Object> {
-    private final @NotNull LaunchOrStopButtonTableCellEditor myEditor;
-
-    @UiThread
-    @VisibleForTesting
-    SetEnabled(@NotNull LaunchOrStopButtonTableCellEditor editor) {
-      myEditor = editor;
-    }
-
-    @UiThread
-    @Override
-    public void onSuccess(@Nullable Object result) {
-      myEditor.myButton.setEnabled(true);
-      myEditor.fireEditingCanceled();
-    }
-
-    @UiThread
-    @Override
-    public void onFailure(@NotNull Throwable throwable) {
-      myEditor.myButton.setEnabled(true);
-      myEditor.fireEditingCanceled();
-
-      myEditor.myShowErrorDialog.accept(throwable, myEditor.myProject);
-    }
+  @NotNull Object getDevice() {
+    return myDevice;
   }
 
-  @UiThread
   @Override
   public @NotNull Component getTableCellEditorComponent(@NotNull JTable table,
                                                         @NotNull Object value,
@@ -175,19 +49,12 @@ final class LaunchOrStopButtonTableCellEditor extends IconButtonTableCellEditor
                                                         int viewRowIndex,
                                                         int viewColumnIndex) {
     myDevice = ((VirtualDeviceTable)table).getDeviceAt(viewRowIndex);
+    VirtualDevice.State state = (VirtualDevice.State)value;
 
-    if (myDevice.isOnline()) {
-      myButton.setDefaultIcon(StudioIcons.Avd.STOP);
-      myButton.setEnabled(true);
-      myButton.setToolTipText("Stop the emulator running this AVD");
-    }
-    else {
-      myButton.setDefaultIcon(StudioIcons.Avd.RUN);
-      myButton.setEnabled(myDevice.getAvdInfo().getStatus().equals(AvdStatus.OK));
-      myButton.setToolTipText("Launch this AVD in the emulator");
-    }
+    myButton.setDefaultIcon(state.getIcon());
+    myButton.setEnabled(state.isEnabled(myDevice));
+    myButton.setToolTipText(state.getTooltipText());
 
-    super.getTableCellEditorComponent(table, value, selected, viewRowIndex, viewColumnIndex);
-    return myButton;
+    return super.getTableCellEditorComponent(table, value, selected, viewRowIndex, viewColumnIndex);
   }
 }
diff --git a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellRenderer.java b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellRenderer.java
index c54f8df8b96..e365a819763 100644
--- a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellRenderer.java
+++ b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellRenderer.java
@@ -15,9 +15,7 @@
  */
 package com.android.tools.idea.devicemanager.virtualtab;
 
-import com.android.sdklib.internal.avd.AvdInfo.AvdStatus;
 import com.android.tools.idea.devicemanager.IconButtonTableCellRenderer;
-import icons.StudioIcons;
 import java.awt.Component;
 import javax.swing.JTable;
 import org.jetbrains.annotations.NotNull;
@@ -30,20 +28,12 @@ final class LaunchOrStopButtonTableCellRenderer extends IconButtonTableCellRende
                                                           boolean focused,
                                                           int viewRowIndex,
                                                           int viewColumnIndex) {
-    VirtualDevice device = ((VirtualDeviceTable)table).getDeviceAt(viewRowIndex);
+    VirtualDevice.State state = (VirtualDevice.State)value;
 
-    if (device.isOnline()) {
-      myButton.setDefaultIcon(StudioIcons.Avd.STOP);
-      myButton.setEnabled(true);
-      myButton.setToolTipText("Stop the emulator running this AVD");
-    }
-    else {
-      myButton.setDefaultIcon(StudioIcons.Avd.RUN);
-      myButton.setEnabled(device.getAvdInfo().getStatus().equals(AvdStatus.OK));
-      myButton.setToolTipText("Launch this AVD in the emulator");
-    }
+    myButton.setDefaultIcon(state.getIcon());
+    myButton.setEnabled(state.isEnabled(((VirtualDeviceTable)table).getDeviceAt(viewRowIndex)));
+    myButton.setToolTipText(state.getTooltipText());
 
-    super.getTableCellRendererComponent(table, value, selected, focused, viewRowIndex, viewColumnIndex);
-    return myButton;
+    return super.getTableCellRendererComponent(table, value, selected, focused, viewRowIndex, viewColumnIndex);
   }
 }
diff --git a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDevice.java b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDevice.java
index 7f2e390fa1f..bae23d81d0a 100644
--- a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDevice.java
+++ b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDevice.java
@@ -117,6 +117,7 @@ public final class VirtualDevice extends Device {
   enum State {
     STOPPED(false, StudioIcons.Avd.RUN, "Launch this AVD in the emulator") {
       @Override
+      @SuppressWarnings("unused")
       boolean isEnabled(@NotNull VirtualDevice device) {
         return device.myAvdInfo.getStatus().equals(AvdStatus.OK);
       }
@@ -124,6 +125,7 @@ public final class VirtualDevice extends Device {
 
     LAUNCHING(false, StudioIcons.Avd.RUN, "Launch this AVD in the emulator") {
       @Override
+      @SuppressWarnings("unused")
       boolean isEnabled(@NotNull VirtualDevice device) {
         return false;
       }
@@ -131,6 +133,7 @@ public final class VirtualDevice extends Device {
 
     LAUNCHED(true, StudioIcons.Avd.STOP, "Stop the emulator running this AVD") {
       @Override
+      @SuppressWarnings("unused")
       boolean isEnabled(@NotNull VirtualDevice device) {
         return true;
       }
@@ -138,6 +141,7 @@ public final class VirtualDevice extends Device {
 
     STOPPING(true, StudioIcons.Avd.STOP, "Stop the emulator running this AVD") {
       @Override
+      @SuppressWarnings("unused")
       boolean isEnabled(@NotNull VirtualDevice device) {
         return false;
       }
@@ -227,6 +231,10 @@ public final class VirtualDevice extends Device {
     return mySizeOnDisk;
   }
 
+  @NotNull Object getState() {
+    return myState;
+  }
+
   boolean isPairable() {
     return newPairingState().myPairable;
   }
diff --git a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDeviceTable.java b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDeviceTable.java
index a2cc824160c..9b9d12db41d 100644
--- a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDeviceTable.java
+++ b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDeviceTable.java
@@ -89,7 +89,7 @@ public final class VirtualDeviceTable extends DeviceTable<VirtualDevice> impleme
 
     initListener();
 
-    setDefaultEditor(VirtualDevice.State.class, new LaunchOrStopButtonTableCellEditor(project));
+    setDefaultEditor(VirtualDevice.State.class, new LaunchOrStopButtonTableCellEditor());
 
     setDefaultEditor(ActivateDeviceFileExplorerWindowValue.class,
                      new ActivateDeviceFileExplorerWindowButtonTableCellEditor<>(project,
diff --git a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDeviceTableModel.java b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDeviceTableModel.java
index ae35e45f5d1..e445ed67b8d 100644
--- a/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDeviceTableModel.java
+++ b/device-manager/src/com/android/tools/idea/devicemanager/virtualtab/VirtualDeviceTableModel.java
@@ -282,7 +282,7 @@ final class VirtualDeviceTableModel extends AbstractTableModel {
       case SIZE_ON_DISK_MODEL_COLUMN_INDEX:
         return myDevices.get(modelRowIndex).getSizeOnDisk();
       case LAUNCH_OR_STOP_MODEL_COLUMN_INDEX:
-        return VirtualDevice.State.STOPPED;
+        return myDevices.get(modelRowIndex).getState();
       case ACTIVATE_DEVICE_FILE_EXPLORER_WINDOW_MODEL_COLUMN_INDEX:
         return ActivateDeviceFileExplorerWindowValue.INSTANCE;
       case EDIT_MODEL_COLUMN_INDEX:
@@ -319,13 +319,13 @@ final class VirtualDeviceTableModel extends AbstractTableModel {
   private void launchOrStop(@NotNull VirtualDevice.State state, int modelRowIndex) {
     switch (state) {
       case STOPPED:
-        // TODO Uncomment the assert when http://b/237442318 is fixed
-        // assert false;
+        assert false;
         break;
       case LAUNCHING:
         launch(modelRowIndex);
         break;
       case LAUNCHED:
+        // noinspection DuplicateBranchesInSwitch
         assert false;
         break;
       case STOPPING:
diff --git a/device-manager/testSrc/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellEditorTest.java b/device-manager/testSrc/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellEditorTest.java
index 0a35dc2689d..0dcdb48ca34 100644
--- a/device-manager/testSrc/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellEditorTest.java
+++ b/device-manager/testSrc/com/android/tools/idea/devicemanager/virtualtab/LaunchOrStopButtonTableCellEditorTest.java
@@ -16,233 +16,77 @@
 package com.android.tools.idea.devicemanager.virtualtab;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
-import com.android.ddmlib.EmulatorConsole;
-import com.android.ddmlib.IDevice;
-import com.android.sdklib.AndroidVersion;
 import com.android.sdklib.internal.avd.AvdInfo;
 import com.android.sdklib.internal.avd.AvdInfo.AvdStatus;
-import com.android.tools.idea.avdmanager.AvdManagerConnection;
-import com.android.tools.idea.devicemanager.CountDownLatchAssert;
-import com.android.tools.idea.devicemanager.CountDownLatchFutureCallback;
-import com.android.tools.idea.devicemanager.DeviceManagerAndroidDebugBridge;
 import com.android.tools.idea.devicemanager.DeviceTables;
 import com.android.tools.idea.devicemanager.IconButton;
-import com.android.tools.idea.devicemanager.IconButtonTableCellEditor;
-import com.android.tools.idea.devicemanager.Key;
-import com.android.tools.idea.devicemanager.virtualtab.LaunchOrStopButtonTableCellEditor.SetEnabled;
-import com.google.common.util.concurrent.FutureCallback;
-import com.google.common.util.concurrent.Futures;
-import com.intellij.openapi.project.Project;
 import icons.StudioIcons;
 import java.util.Optional;
-import java.util.concurrent.CountDownLatch;
-import java.util.function.BiConsumer;
-import javax.swing.AbstractButton;
 import javax.swing.JTable;
 import javax.swing.event.CellEditorListener;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
-import org.mockito.ArgumentMatcher;
-import org.mockito.ArgumentMatchers;
 import org.mockito.Mockito;
 
 @RunWith(JUnit4.class)
 public final class LaunchOrStopButtonTableCellEditorTest {
-  private final @NotNull AvdManagerConnection myConnection = Mockito.mock(AvdManagerConnection.class);
-  private final @NotNull AvdInfo myAvd = Mockito.mock(AvdInfo.class);
+  private final @NotNull AvdInfo myAvd;
+  private final @NotNull CellEditorListener myListener;
+  private final @NotNull LaunchOrStopButtonTableCellEditor myEditor;
 
-  private IconButtonTableCellEditor myEditor;
+  public LaunchOrStopButtonTableCellEditorTest() {
+    myAvd = Mockito.mock(AvdInfo.class);
+    myListener = Mockito.mock(CellEditorListener.class);
 
-  @Test
-  public void onSuccessDeviceIsOnline() throws InterruptedException {
-    // Arrange
-    Key key = TestVirtualDevices.newKey("Pixel_5_API_31");
-
-    DeviceManagerAndroidDebugBridge bridge = Mockito.mock(DeviceManagerAndroidDebugBridge.class);
-    Mockito.when(bridge.findDevice(null, key)).thenReturn(Futures.immediateFuture(Mockito.mock(IDevice.class)));
-
-    CountDownLatch latch = new CountDownLatch(1);
-    EmulatorConsole console = Mockito.mock(EmulatorConsole.class);
-    CellEditorListener listener = Mockito.mock(CellEditorListener.class);
-    JTable table = DeviceTables.mock(TestVirtualDevices.onlinePixel5Api31(myAvd));
-
-    myEditor = new LaunchOrStopButtonTableCellEditor(null,
-                                                     bridge,
-                                                     editor -> newSetEnabled(editor, latch),
-                                                     device -> console,
-                                                     AvdManagerConnection::getDefaultAvdManagerConnection,
-                                                     VirtualTabMessages::showErrorDialog);
-
-    myEditor.addCellEditorListener(listener);
-    myEditor.getTableCellEditorComponent(table, VirtualDevice.State.STOPPED, false, 0, 3);
-
-    AbstractButton button = myEditor.getButton();
-
-    // Act
-    button.doClick();
-
-    // Assert
-    CountDownLatchAssert.await(latch);
-
-    Mockito.verify(console).kill();
-    assertTrue(button.isEnabled());
-    Mockito.verify(listener).editingCanceled(myEditor.getChangeEvent());
-  }
-
-  @Test
-  public void onSuccess() throws InterruptedException {
-    // Arrange
-    CountDownLatch latch = new CountDownLatch(1);
-    Mockito.when(myConnection.startAvd(null, myAvd)).thenReturn(Futures.immediateFuture(Mockito.mock(IDevice.class)));
-    CellEditorListener listener = Mockito.mock(CellEditorListener.class);
-    Mockito.when(myAvd.getStatus()).thenReturn(AvdStatus.OK);
-    JTable table = DeviceTables.mock(TestVirtualDevices.pixel5Api31(myAvd));
-
-    myEditor = new LaunchOrStopButtonTableCellEditor(null,
-                                                     new DeviceManagerAndroidDebugBridge(),
-                                                     editor -> newSetEnabled(editor, latch),
-                                                     EmulatorConsole::getConsole,
-                                                     () -> myConnection,
-                                                     VirtualTabMessages::showErrorDialog);
-
-    myEditor.addCellEditorListener(listener);
-    myEditor.getTableCellEditorComponent(table, VirtualDevice.State.STOPPED, false, 0, 3);
-
-    AbstractButton button = myEditor.getButton();
-
-    // Act
-    button.doClick();
-
-    // Assert
-    CountDownLatchAssert.await(latch);
-
-    assertTrue(button.isEnabled());
-    Mockito.verify(listener).editingCanceled(myEditor.getChangeEvent());
-  }
-
-  @Test
-  public void onFailure() throws Exception {
-    // Arrange
-    DeviceManagerAndroidDebugBridge bridge = Mockito.mock(DeviceManagerAndroidDebugBridge.class);
-    Mockito.when(bridge.findDevice(null, TestVirtualDevices.newKey("Pixel_5_API_31"))).thenReturn(Futures.immediateFuture(null));
-
-    CountDownLatch latch = new CountDownLatch(1);
-
-    @SuppressWarnings("unchecked")
-    BiConsumer<Throwable, Project> showErrorDialog = Mockito.mock(BiConsumer.class);
-
-    CellEditorListener listener = Mockito.mock(CellEditorListener.class);
-    JTable table = DeviceTables.mock(TestVirtualDevices.onlinePixel5Api31(myAvd));
-
-    myEditor = new LaunchOrStopButtonTableCellEditor(null,
-                                                     bridge,
-                                                     editor -> newSetEnabled(editor, latch),
-                                                     EmulatorConsole::getConsole,
-                                                     AvdManagerConnection::getDefaultAvdManagerConnection,
-                                                     showErrorDialog);
-
-    myEditor.addCellEditorListener(listener);
-    myEditor.getTableCellEditorComponent(table, VirtualDevice.State.STOPPED, false, 0, 3);
-
-    AbstractButton button = myEditor.getButton();
-
-    // Act
-    button.doClick();
-
-    // Assert
-    CountDownLatchAssert.await(latch);
-
-    assertTrue(button.isEnabled());
-    Mockito.verify(listener).editingCanceled(myEditor.getChangeEvent());
-
-    ArgumentMatcher<Throwable> matcher = actual -> matches("Unable to stop Pixel 5 API 31",
-                                                           "An error occurred stopping Pixel 5 API 31. To stop the device, try manually " +
-                                                           "closing the Pixel 5 API 31 emulator window.",
-                                                           actual);
-
-    Mockito.verify(showErrorDialog).accept(ArgumentMatchers.argThat(matcher), ArgumentMatchers.isNull());
-  }
-
-  private static boolean matches(@NotNull String expectedTitle, @NotNull String expectedMessage, @NotNull Throwable actual) {
-    if (!(actual instanceof ErrorDialogException)) {
-      return false;
-    }
-
-    ErrorDialogException exception = ((ErrorDialogException)actual);
-    return expectedTitle.equals(exception.getTitle()) && expectedMessage.equals(exception.getMessage());
-  }
-
-  private static @NotNull FutureCallback<@Nullable Object> newSetEnabled(@NotNull LaunchOrStopButtonTableCellEditor editor,
-                                                                         @NotNull CountDownLatch latch) {
-    return new CountDownLatchFutureCallback<>(new SetEnabled(editor), latch);
+    myEditor = new LaunchOrStopButtonTableCellEditor();
+    myEditor.addCellEditorListener(myListener);
   }
 
   @Test
   public void getTableCellEditorComponentDeviceIsOnline() {
     // Arrange
-    myEditor = new LaunchOrStopButtonTableCellEditor(null);
-    JTable table = DeviceTables.mock(TestVirtualDevices.onlinePixel5Api31(myAvd));
-
-    // Act
-    Object component = myEditor.getTableCellEditorComponent(table, VirtualDevice.State.STOPPED, false, 0, 3);
-
-    // Assert
-    IconButton button = (IconButton)component;
-
-    assertEquals(Optional.of(StudioIcons.Avd.STOP), button.getDefaultIcon());
-    assertTrue(button.isEnabled());
-    assertEquals("Stop the emulator running this AVD", button.getToolTipText());
-  }
-
-  @Test
-  public void getTableCellEditorComponentStatusDoesntEqualOk() {
-    // Arrange
-    myEditor = new LaunchOrStopButtonTableCellEditor(null);
-    Mockito.when(myAvd.getStatus()).thenReturn(AvdStatus.ERROR_PROPERTIES);
-
-    VirtualDevice device = new VirtualDevice.Builder()
-      .setKey(TestVirtualDevices.newKey("Pixel_5_API_31"))
-      .setName("Pixel 5 API 31")
-      .setTarget("Android 12.0")
-      .setCpuArchitecture("arm")
-      .setAndroidVersion(new AndroidVersion(31))
-      .setAvdInfo(myAvd)
-      .build();
-
+    VirtualDevice device = TestVirtualDevices.onlinePixel5Api31(myAvd);
     JTable table = DeviceTables.mock(device);
 
     // Act
-    Object component = myEditor.getTableCellEditorComponent(table, VirtualDevice.State.STOPPED, false, 0, 3);
+    IconButton component = (IconButton)myEditor.getTableCellEditorComponent(table, VirtualDevice.State.LAUNCHED, false, 0, 3);
+    component.doClick();
 
     // Assert
-    IconButton button = (IconButton)component;
+    assertEquals(device, myEditor.getDevice());
+
+    assertEquals(Optional.of(StudioIcons.Avd.STOP), component.getDefaultIcon());
+    assertTrue(component.isEnabled());
+    assertEquals("Stop the emulator running this AVD", component.getToolTipText());
 
-    assertEquals(Optional.of(StudioIcons.Avd.RUN), button.getDefaultIcon());
-    assertFalse(button.isEnabled());
-    assertEquals("Launch this AVD in the emulator", button.getToolTipText());
+    assertEquals(VirtualDevice.State.STOPPING, myEditor.getCellEditorValue());
+    Mockito.verify(myListener).editingStopped(myEditor.getChangeEvent());
   }
 
   @Test
   public void getTableCellEditorComponent() {
     // Arrange
-    myEditor = new LaunchOrStopButtonTableCellEditor(null);
     Mockito.when(myAvd.getStatus()).thenReturn(AvdStatus.OK);
-    JTable table = DeviceTables.mock(TestVirtualDevices.pixel5Api31(myAvd));
+
+    VirtualDevice device = TestVirtualDevices.pixel5Api31(myAvd);
+    JTable table = DeviceTables.mock(device);
 
     // Act
-    Object component = myEditor.getTableCellEditorComponent(table, VirtualDevice.State.STOPPED, false, 0, 3);
+    IconButton component = (IconButton)myEditor.getTableCellEditorComponent(table, VirtualDevice.State.STOPPED, false, 0, 3);
+    component.doClick();
 
     // Assert
-    IconButton button = (IconButton)component;
+    assertEquals(device, myEditor.getDevice());
+
+    assertEquals(Optional.of(StudioIcons.Avd.RUN), component.getDefaultIcon());
+    assertTrue(component.isEnabled());
+    assertEquals("Launch this AVD in the emulator", component.getToolTipText());
 
-    assertEquals(Optional.of(StudioIcons.Avd.RUN), button.getDefaultIcon());
-    assertTrue(button.isEnabled());
-    assertEquals("Launch this AVD in the emulator", button.getToolTipText());
+    assertEquals(VirtualDevice.State.LAUNCHING, myEditor.getCellEditorValue());
+    Mockito.verify(myListener).editingStopped(myEditor.getChangeEvent());
   }
 }