Integrate uiautomatorviewer into Eclipse

This CL integrates uiautomatorviewer into Eclipse. There are two
components to this:
 1. Take UI Hierarchy Dump from a device: This is added as an action
    to the device view in the DDMS perspective.
 2. Viewing the resulting UI hierarchy: This is added as an editor,
    associated with the .uix extension. The editor simply wraps a
    UiAutomatorView.

The UiAutomatorView is modified slightly to allow for opening just
the hierarchy file without the associated screenshot. This is to
enable opening the .uix file via standard Eclipse Open File Dialog.
When opened this way, the canvas displaying the screenshot is replaced
with a button that allows the user to specify the actual screenshot
to be displayed.

Change-Id: I0f31fe95b2728123ff6f0f748478c18c0c2f1b48

6 files changed, 304 insertions, 151 deletions

diff --git a/uiautomatorviewer/src/com/android/uiautomator/OpenDialog.java b/uiautomatorviewer/src/com/android/uiautomator/OpenDialog.java
index 677d2f7..97a437b 100644
--- a/uiautomatorviewer/src/com/android/uiautomator/OpenDialog.java
+++ b/uiautomatorviewer/src/com/android/uiautomator/OpenDialog.java
@@ -196,7 +196,7 @@ public class OpenDialog extends Dialog {
         if (initialPath.isEmpty() && sScreenshotFile != null && sScreenshotFile.isFile()) {
             initialPath = sScreenshotFile.getParentFile().getAbsolutePath();
         }
-        String[] filter = {"*.xml"};
+        String[] filter = {"*.uix"};
         fd.setFilterExtensions(filter);
         String selected = fd.open();
         if (selected != null) {
@@ -208,8 +208,7 @@ public class OpenDialog extends Dialog {
     }
 
     private void updateButtonState() {
-        mOkButton.setEnabled(sScreenshotFile != null && sXmlDumpFile != null
-                && sScreenshotFile.isFile() && sXmlDumpFile.isFile());
+        mOkButton.setEnabled(sXmlDumpFile != null && sXmlDumpFile.isFile());
     }
 
     public boolean hasFileChanged() {
diff --git a/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorHelper.java b/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorHelper.java
new file mode 100644
index 0000000..7571c8b
--- /dev/null
+++ b/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorHelper.java
@@ -0,0 +1,184 @@
+/*
+ * Copyright (C) 2012 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.uiautomator;
+
+import com.android.ddmlib.CollectingOutputReceiver;
+import com.android.ddmlib.IDevice;
+import com.android.ddmlib.RawImage;
+import com.android.ddmlib.SyncService;
+
+import org.eclipse.core.runtime.IProgressMonitor;
+import org.eclipse.core.runtime.NullProgressMonitor;
+import org.eclipse.swt.SWT;
+import org.eclipse.swt.graphics.Image;
+import org.eclipse.swt.graphics.ImageData;
+import org.eclipse.swt.graphics.ImageLoader;
+import org.eclipse.swt.graphics.PaletteData;
+import org.eclipse.swt.widgets.Display;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+public class UiAutomatorHelper {
+    public static final int UIAUTOMATOR_MIN_API_LEVEL = 16;
+
+    private static final String UIAUTOMATOR = "/system/bin/uiautomator";    //$NON-NLS-1$
+    private static final String UIAUTOMATOR_DUMP_COMMAND = "dump";          //$NON-NLS-1$
+    private static final String UIDUMP_DEVICE_PATH = "/sdcard/uidump.xml";  //$NON-NLS-1$
+
+    private static boolean supportsUiAutomator(IDevice device) {
+        String apiLevelString = device.getProperty(IDevice.PROP_BUILD_API_LEVEL);
+        int apiLevel;
+        try {
+            apiLevel = Integer.parseInt(apiLevelString);
+        } catch (NumberFormatException e) {
+            apiLevel = UIAUTOMATOR_MIN_API_LEVEL;
+        }
+
+        return apiLevel >= UIAUTOMATOR_MIN_API_LEVEL;
+    }
+
+    private static void getUiHierarchyFile(IDevice device, File dst, IProgressMonitor monitor) {
+        if (monitor == null) {
+            monitor = new NullProgressMonitor();
+        }
+
+        monitor.subTask("Deleting old UI XML snapshot ...");
+        String command = "rm " + UIDUMP_DEVICE_PATH;
+
+        try {
+            CountDownLatch commandCompleteLatch = new CountDownLatch(1);
+            device.executeShellCommand(command,
+                    new CollectingOutputReceiver(commandCompleteLatch));
+            commandCompleteLatch.await(5, TimeUnit.SECONDS);
+        } catch (Exception e1) {
+            // ignore exceptions while deleting stale files
+        }
+
+        monitor.subTask("Taking UI XML snapshot...");
+        command = String.format("%s %s %s", UIAUTOMATOR,
+                UIAUTOMATOR_DUMP_COMMAND,
+                UIDUMP_DEVICE_PATH);
+        CountDownLatch commandCompleteLatch = new CountDownLatch(1);
+
+        try {
+            device.executeShellCommand(command,
+                    new CollectingOutputReceiver(commandCompleteLatch));
+            commandCompleteLatch.await(20, TimeUnit.SECONDS);
+
+            monitor.subTask("Pull UI XML snapshot from device...");
+            device.getSyncService().pullFile(UIDUMP_DEVICE_PATH,
+                    dst.getAbsolutePath(), SyncService.getNullProgressMonitor());
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public static UiAutomatorResult takeSnapshot(IDevice device, IProgressMonitor monitor)
+                                throws UiAutomatorException {
+        if (monitor == null) {
+            monitor = new NullProgressMonitor();
+        }
+
+        monitor.subTask("Checking if device support UI Automator");
+        if (!supportsUiAutomator(device)) {
+            String msg = "UI Automator requires a device with API Level "
+                                + UIAUTOMATOR_MIN_API_LEVEL;
+            throw new UiAutomatorException(msg, null);
+        }
+
+        monitor.subTask("Creating temporary files for uiautomator results.");
+        File tmpDir = null;
+        File xmlDumpFile = null;
+        File screenshotFile = null;
+        try {
+            tmpDir = File.createTempFile("uiautomatorviewer_", "");
+            tmpDir.delete();
+            if (!tmpDir.mkdirs())
+                throw new IOException("Failed to mkdir");
+            xmlDumpFile = File.createTempFile("dump_", ".uix", tmpDir);
+            screenshotFile = File.createTempFile("screenshot_", ".png", tmpDir);
+        } catch (Exception e) {
+            String msg = "Error while creating temporary file to save snapshot: "
+                    + e.getMessage();
+            throw new UiAutomatorException(msg, e);
+        }
+
+        tmpDir.deleteOnExit();
+        xmlDumpFile.deleteOnExit();
+        screenshotFile.deleteOnExit();
+
+        monitor.subTask("Obtaining UI hierarchy");
+        try {
+            UiAutomatorHelper.getUiHierarchyFile(device, xmlDumpFile, monitor);
+        } catch (Exception e) {
+            String msg = "Error while obtaining UI hierarchy XML file: " + e.getMessage();
+            throw new UiAutomatorException(msg, e);
+        }
+
+        UiAutomatorModel model;
+        try {
+            model = new UiAutomatorModel(xmlDumpFile);
+        } catch (Exception e) {
+            String msg = "Error while parsing UI hierarchy XML file: " + e.getMessage();
+            throw new UiAutomatorException(msg, e);
+        }
+
+        monitor.subTask("Obtaining device screenshot");
+        RawImage rawImage;
+        try {
+            rawImage = device.getScreenshot();
+        } catch (Exception e) {
+            String msg = "Error taking device screenshot: " + e.getMessage();
+            throw new UiAutomatorException(msg, e);
+        }
+
+        PaletteData palette = new PaletteData(
+                rawImage.getRedMask(),
+                rawImage.getGreenMask(),
+                rawImage.getBlueMask());
+        ImageData imageData = new ImageData(rawImage.width, rawImage.height,
+                rawImage.bpp, palette, 1, rawImage.data);
+        ImageLoader loader = new ImageLoader();
+        loader.data = new ImageData[] { imageData };
+        loader.save(screenshotFile.getAbsolutePath(), SWT.IMAGE_PNG);
+        Image screenshot = new Image(Display.getDefault(), imageData);
+
+        return new UiAutomatorResult(xmlDumpFile, model, screenshot);
+    }
+
+    @SuppressWarnings("serial")
+    public static class UiAutomatorException extends Exception {
+        public UiAutomatorException(String msg, Throwable t) {
+            super(msg, t);
+        }
+    }
+
+    public static class UiAutomatorResult {
+        public final File uiHierarchy;
+        public final UiAutomatorModel model;
+        public final Image screenshot;
+
+        public UiAutomatorResult(File uiXml, UiAutomatorModel m, Image s) {
+            uiHierarchy = uiXml;
+            model = m;
+            screenshot = s;
+        }
+    }
+}
diff --git a/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorView.java b/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorView.java
index 6e943c6..c5f3e59 100644
--- a/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorView.java
+++ b/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorView.java
@@ -40,28 +40,45 @@ import org.eclipse.jface.viewers.TextCellEditor;
 import org.eclipse.jface.viewers.TreeViewer;
 import org.eclipse.swt.SWT;
 import org.eclipse.swt.custom.SashForm;
+import org.eclipse.swt.custom.StackLayout;
 import org.eclipse.swt.events.MouseAdapter;
 import org.eclipse.swt.events.MouseEvent;
 import org.eclipse.swt.events.MouseMoveListener;
 import org.eclipse.swt.events.PaintEvent;
 import org.eclipse.swt.events.PaintListener;
+import org.eclipse.swt.events.SelectionAdapter;
+import org.eclipse.swt.events.SelectionEvent;
 import org.eclipse.swt.graphics.Image;
+import org.eclipse.swt.graphics.ImageData;
+import org.eclipse.swt.graphics.ImageLoader;
 import org.eclipse.swt.graphics.Rectangle;
 import org.eclipse.swt.graphics.Transform;
 import org.eclipse.swt.layout.FillLayout;
 import org.eclipse.swt.layout.GridData;
 import org.eclipse.swt.layout.GridLayout;
+import org.eclipse.swt.widgets.Button;
 import org.eclipse.swt.widgets.Canvas;
 import org.eclipse.swt.widgets.Composite;
+import org.eclipse.swt.widgets.Display;
+import org.eclipse.swt.widgets.FileDialog;
 import org.eclipse.swt.widgets.Group;
 import org.eclipse.swt.widgets.Table;
 import org.eclipse.swt.widgets.TableColumn;
 import org.eclipse.swt.widgets.Tree;
 
+import java.io.File;
+
 public class UiAutomatorView extends Composite {
     private static final int IMG_BORDER = 2;
 
+    // The screenshot area is made of a stack layout of two components: screenshot canvas and
+    // a "specify screenshot" button. If a screenshot is already available, then that is displayed
+    // on the canvas. If it is not availble, then the "specify screenshot" button is displayed.
+    private Composite mScreenshotComposite;
+    private StackLayout mStackLayout;
+    private Composite mSetScreenshotComposite;
     private Canvas mScreenshotCanvas;
+
     private TreeViewer mTreeViewer;
     private TableViewer mTableViewer;
 
@@ -69,6 +86,7 @@ public class UiAutomatorView extends Composite {
     private int mDx, mDy;
 
     private UiAutomatorModel mModel;
+    private File mModelFile;
     private Image mScreenshot;
 
     public UiAutomatorView(Composite parent, int style) {
@@ -76,8 +94,16 @@ public class UiAutomatorView extends Composite {
         setLayout(new FillLayout());
 
         SashForm baseSash = new SashForm(this, SWT.HORIZONTAL);
+
+        mScreenshotComposite = new Composite(baseSash, SWT.BORDER);
+        mStackLayout = new StackLayout();
+        mScreenshotComposite.setLayout(mStackLayout);
+
         // draw the canvas with border, so the divider area for sash form can be highlighted
-        mScreenshotCanvas = new Canvas(baseSash, SWT.BORDER);
+        mScreenshotCanvas = new Canvas(mScreenshotComposite, SWT.BORDER);
+        mStackLayout.topControl = mScreenshotCanvas;
+        mScreenshotComposite.layout();
+
         mScreenshotCanvas.addMouseListener(new MouseAdapter() {
             @Override
             public void mouseUp(MouseEvent e) {
@@ -155,6 +181,43 @@ public class UiAutomatorView extends Composite {
             }
         });
 
+        mSetScreenshotComposite = new Composite(mScreenshotComposite, SWT.NONE);
+        mSetScreenshotComposite.setLayout(new GridLayout());
+
+        final Button setScreenshotButton = new Button(mSetScreenshotComposite, SWT.PUSH);
+        setScreenshotButton.setText("Specify Screenshot...");
+        setScreenshotButton.addSelectionListener(new SelectionAdapter() {
+            @Override
+            public void widgetSelected(SelectionEvent arg0) {
+                FileDialog fd = new FileDialog(setScreenshotButton.getShell());
+                fd.setFilterExtensions(new String[] { "*.png" });
+                if (mModelFile != null) {
+                    fd.setFilterPath(mModelFile.getParent());
+                }
+                String screenshotPath = fd.open();
+                if (screenshotPath == null) {
+                    return;
+                }
+
+                ImageData[] data;
+                try {
+                    data = new ImageLoader().load(screenshotPath);
+                } catch (Exception e) {
+                    return;
+                }
+
+                // "data" is an array, probably used to handle images that has multiple frames
+                // i.e. gifs or icons, we just care if it has at least one here
+                if (data.length < 1) {
+                    return;
+                }
+
+                mScreenshot = new Image(Display.getDefault(), data[0]);
+                redrawScreenshot();
+            }
+        });
+
+
         // right sash is split into 2 parts: upper-right and lower-right
         // both are composites with borders, so that the horizontal divider can be highlighted by
         // the borders
@@ -314,6 +377,13 @@ public class UiAutomatorView extends Composite {
      * retrieved from Model
      */
     public void redrawScreenshot() {
+        if (mScreenshot == null) {
+            mStackLayout.topControl = mSetScreenshotComposite;
+        } else {
+            mStackLayout.topControl = mScreenshotCanvas;
+        }
+        mScreenshotComposite.layout();
+
         mScreenshotCanvas.redraw();
     }
 
@@ -334,8 +404,9 @@ public class UiAutomatorView extends Composite {
         mTreeViewer.setSelection(new StructuredSelection(node), true);
     }
 
-    public void setModel(UiAutomatorModel model, Image screenshot) {
+    public void setModel(UiAutomatorModel model, File modelBackingFile, Image screenshot) {
         mModel = model;
+        mModelFile = modelBackingFile;
 
         if (mScreenshot != null) {
             mScreenshot.dispose();
diff --git a/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorViewer.java b/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorViewer.java
index d0f3b94..37018b4 100644
--- a/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorViewer.java
+++ b/uiautomatorviewer/src/com/android/uiautomator/UiAutomatorViewer.java
@@ -32,6 +32,8 @@ import org.eclipse.swt.widgets.Display;
 import org.eclipse.swt.widgets.Shell;
 import org.eclipse.swt.widgets.ToolBar;
 
+import java.io.File;
+
 public class UiAutomatorViewer extends ApplicationWindow {
     private UiAutomatorView mUiAutomatorView;
 
@@ -90,16 +92,17 @@ public class UiAutomatorViewer extends ApplicationWindow {
         return new Point(800, 600);
     }
 
-    public void setModel(final UiAutomatorModel model, final Image screenshot) {
+    public void setModel(final UiAutomatorModel model, final File modelFile,
+                                                                final Image screenshot) {
         if (Display.getDefault().getThread() != Thread.currentThread()) {
             Display.getDefault().syncExec(new Runnable() {
                 @Override
                 public void run() {
-                    mUiAutomatorView.setModel(model, screenshot);
+                    mUiAutomatorView.setModel(model, modelFile, screenshot);
                 }
             });
         } else {
-            mUiAutomatorView.setModel(model, screenshot);
+            mUiAutomatorView.setModel(model, modelFile, screenshot);
         }
     }
 }
diff --git a/uiautomatorviewer/src/com/android/uiautomator/actions/OpenFilesAction.java b/uiautomatorviewer/src/com/android/uiautomator/actions/OpenFilesAction.java
index 0d4e707..46ee9b6 100644
--- a/uiautomatorviewer/src/com/android/uiautomator/actions/OpenFilesAction.java
+++ b/uiautomatorviewer/src/com/android/uiautomator/actions/OpenFilesAction.java
@@ -59,23 +59,25 @@ public class OpenFilesAction extends Action {
         }
 
         Image img = null;
-        try {
-            File screenshot = d.getScreenshotFile();
-            ImageData[] data = new ImageLoader().load(screenshot.getAbsolutePath());
+        File screenshot = d.getScreenshotFile();
+        if (screenshot != null) {
+            try {
+                ImageData[] data = new ImageLoader().load(screenshot.getAbsolutePath());
 
-            // "data" is an array, probably used to handle images that has multiple frames
-            // i.e. gifs or icons, we just care if it has at least one here
-            if (data.length < 1) {
-                throw new RuntimeException("Unable to load image: "
-                                + screenshot.getAbsolutePath());
-            }
+                // "data" is an array, probably used to handle images that has multiple frames
+                // i.e. gifs or icons, we just care if it has at least one here
+                if (data.length < 1) {
+                    throw new RuntimeException("Unable to load image: "
+                            + screenshot.getAbsolutePath());
+                }
 
-            img = new Image(Display.getDefault(), data[0]);
-        } catch (Exception e) {
-            // FIXME: show error
-            return;
+                img = new Image(Display.getDefault(), data[0]);
+            } catch (Exception e) {
+                // FIXME: show error
+                return;
+            }
         }
 
-        mViewer.setModel(model, img);
+        mViewer.setModel(model, d.getXmlDumpFile(), img);
     }
 }
diff --git a/uiautomatorviewer/src/com/android/uiautomator/actions/ScreenshotAction.java b/uiautomatorviewer/src/com/android/uiautomator/actions/ScreenshotAction.java
index ff09779..700b041 100644
--- a/uiautomatorviewer/src/com/android/uiautomator/actions/ScreenshotAction.java
+++ b/uiautomatorviewer/src/com/android/uiautomator/actions/ScreenshotAction.java
@@ -16,12 +16,11 @@
 
 package com.android.uiautomator.actions;
 
-import com.android.ddmlib.CollectingOutputReceiver;
 import com.android.ddmlib.IDevice;
-import com.android.ddmlib.RawImage;
-import com.android.ddmlib.SyncService;
 import com.android.uiautomator.DebugBridge;
-import com.android.uiautomator.UiAutomatorModel;
+import com.android.uiautomator.UiAutomatorHelper;
+import com.android.uiautomator.UiAutomatorHelper.UiAutomatorException;
+import com.android.uiautomator.UiAutomatorHelper.UiAutomatorResult;
 import com.android.uiautomator.UiAutomatorViewer;
 
 import org.eclipse.core.runtime.IProgressMonitor;
@@ -38,32 +37,17 @@ import org.eclipse.jface.window.Window;
 import org.eclipse.swt.SWT;
 import org.eclipse.swt.events.SelectionAdapter;
 import org.eclipse.swt.events.SelectionEvent;
-import org.eclipse.swt.graphics.Image;
-import org.eclipse.swt.graphics.ImageData;
-import org.eclipse.swt.graphics.ImageLoader;
-import org.eclipse.swt.graphics.PaletteData;
 import org.eclipse.swt.layout.GridLayout;
 import org.eclipse.swt.widgets.Combo;
 import org.eclipse.swt.widgets.Composite;
 import org.eclipse.swt.widgets.Control;
-import org.eclipse.swt.widgets.Display;
 import org.eclipse.swt.widgets.Label;
 import org.eclipse.swt.widgets.Shell;
 
-import java.io.File;
-import java.io.IOException;
 import java.lang.reflect.InvocationTargetException;
 import java.util.List;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 
 public class ScreenshotAction extends Action {
-    private static final String UIAUTOMATOR = "/system/bin/uiautomator";    //$NON-NLS-1$
-    private static final String UIAUTOMATOR_DUMP_COMMAND = "dump";          //$NON-NLS-1$
-    private static final String UIDUMP_DEVICE_PATH = "/sdcard/uidump.xml";  //$NON-NLS-1$
-
-    private static final int MIN_API_LEVEL = 16;
-
     UiAutomatorViewer mViewer;
 
     public ScreenshotAction(UiAutomatorViewer viewer) {
@@ -93,128 +77,38 @@ public class ScreenshotAction extends Action {
         ProgressMonitorDialog dialog = new ProgressMonitorDialog(mViewer.getShell());
         try {
             dialog.run(true, false, new IRunnableWithProgress() {
-                private void showError(final String msg, final Throwable t,
-                        IProgressMonitor monitor) {
-                    monitor.done();
-                    mViewer.getShell().getDisplay().syncExec(new Runnable() {
-                        @Override
-                        public void run() {
-                            Status s = new Status(IStatus.ERROR, "Screenshot", msg, t);
-                            ErrorDialog.openError(
-                                    mViewer.getShell(), "Error", "Error obtaining UI hierarchy", s);
-                        }
-                    });
-                }
-
                 @Override
                 public void run(IProgressMonitor monitor) throws InvocationTargetException,
-                InterruptedException {
-                    File tmpDir = null;
-                    File xmlDumpFile = null;
-                    File screenshotFile = null;
-                    try {
-                        tmpDir = File.createTempFile("uiautomatorviewer_", "");
-                        tmpDir.delete();
-                        if (!tmpDir.mkdirs())
-                            throw new IOException("Failed to mkdir");
-                        xmlDumpFile = File.createTempFile("dump_", ".xml", tmpDir);
-                        screenshotFile = File.createTempFile("screenshot_", ".png", tmpDir);
-                    } catch (IOException e) {
-                        e.printStackTrace();
-                        showError("Cannot get temp directory", e, monitor);
-                        return;
-                    }
-
-                    tmpDir.deleteOnExit();
-                    xmlDumpFile.deleteOnExit();
-                    screenshotFile.deleteOnExit();
-
-                    String apiLevelString = device.getProperty(IDevice.PROP_BUILD_API_LEVEL);
-                    int apiLevel;
-                    try {
-                        apiLevel = Integer.parseInt(apiLevelString);
-                    } catch (NumberFormatException e) {
-                        apiLevel = MIN_API_LEVEL;
-                    }
-                    if (apiLevel < MIN_API_LEVEL) {
-                        showError("uiautomator requires a device with API Level " + MIN_API_LEVEL,
-                                null, monitor);
-                        return;
-                    }
-
-                    monitor.subTask("Deleting old UI XML snapshot ...");
-                    String command = "rm " + UIDUMP_DEVICE_PATH;
-                    try {
-                        CountDownLatch commandCompleteLatch = new CountDownLatch(1);
-                        device.executeShellCommand(command,
-                                new CollectingOutputReceiver(commandCompleteLatch));
-                        commandCompleteLatch.await(5, TimeUnit.SECONDS);
-                    } catch (Exception e1) {
-                        // ignore exceptions while deleting stale files
-                    }
-
-                    monitor.subTask("Taking UI XML snapshot...");
-                    command = String.format("%s %s %s", UIAUTOMATOR,
-                                                        UIAUTOMATOR_DUMP_COMMAND,
-                                                        UIDUMP_DEVICE_PATH);
-                    try {
-                        CountDownLatch commandCompleteLatch = new CountDownLatch(1);
-                        device.executeShellCommand(command,
-                                new CollectingOutputReceiver(commandCompleteLatch));
-                        commandCompleteLatch.await(5, TimeUnit.SECONDS);
-                    } catch (Exception e1) {
-                        showError("", e1, monitor);
-                        return;
-                    }
-
-                    monitor.subTask("Pull UI XML snapshot from device...");
-                    try {
-                        device.getSyncService().pullFile(UIDUMP_DEVICE_PATH,
-                                xmlDumpFile.getAbsolutePath(), SyncService.getNullProgressMonitor());
-                    } catch (Exception e1) {
-                        showError("Error copying UI XML file from device", e1, monitor);
-                        return;
-                    }
-
-                    monitor.subTask("Taking screenshot...");
-                    RawImage rawImage;
+                                                                        InterruptedException {
+                    UiAutomatorResult result = null;
                     try {
-                        rawImage = device.getScreenshot();
-                    } catch (Exception e1) {
-                        showError("Error taking device screenshot", e1, monitor);
+                        result = UiAutomatorHelper.takeSnapshot(device, monitor);
+                    } catch (UiAutomatorException e) {
+                        monitor.done();
+                        showError(e.getMessage(), e);
                         return;
                     }
 
-                    UiAutomatorModel model;
-                    try {
-                        model = new UiAutomatorModel(xmlDumpFile);
-                    } catch (Exception e) {
-                        showError("Error while parsing UI hierarchy XML file", e, monitor);
-                        return;
-                    }
-
-                    PaletteData palette = new PaletteData(
-                            rawImage.getRedMask(),
-                            rawImage.getGreenMask(),
-                            rawImage.getBlueMask());
-                    ImageData imageData = new ImageData(rawImage.width, rawImage.height,
-                            rawImage.bpp, palette, 1, rawImage.data);
-                    ImageLoader loader = new ImageLoader();
-                    loader.data = new ImageData[] { imageData };
-                    loader.save(screenshotFile.getAbsolutePath(), SWT.IMAGE_PNG);
-                    Image screenshot = new Image(Display.getDefault(), imageData);
-
-                    mViewer.setModel(model, screenshot);
+                    mViewer.setModel(result.model, result.uiHierarchy, result.screenshot);
                     monitor.done();
                 }
             });
-        } catch (InvocationTargetException e) {
-            e.printStackTrace();
-        } catch (InterruptedException e) {
-            e.printStackTrace();
+        } catch (Exception e) {
+            showError("Unexpected error while obtaining UI hierarchy", e);
         }
     }
 
+    private void showError(final String msg, final Throwable t) {
+        mViewer.getShell().getDisplay().syncExec(new Runnable() {
+            @Override
+            public void run() {
+                Status s = new Status(IStatus.ERROR, "Screenshot", msg, t);
+                ErrorDialog.openError(
+                        mViewer.getShell(), "Error", "Error obtaining UI hierarchy", s);
+            }
+        });
+    }
+
     private IDevice pickDevice() {
         List<IDevice> devices = DebugBridge.getDevices();
         if (devices.size() == 0) {