clear internal data while refreshing rootandroid-sdk-4.4.2_r1.0.1android-sdk-4.4.2_r1android-cts-4.4_r4android-cts-4.4_r1android-4.4w_r1android-4.4_r1.2.0.1android-4.4_r1.2android-4.4_r1.1.0.1android-4.4_r1.1android-4.4_r1.0.1android-4.4_r1android-4.4_r0.9android-4.4_r0.8android-4.4_r0.7android-4.4.4_r2.0.1android-4.4.4_r2android-4.4.4_r1.0.1android-4.4.4_r1android-4.4.3_r1.1.0.1android-4.4.3_r1.1android-4.4.3_r1.0.1android-4.4.3_r1android-4.4.2_r2.0.1android-4.4.2_r2android-4.4.2_r1.0.1android-4.4.2_r1android-4.4.1_r1.0.1android-4.4.1_r1kitkat-wearkitkat-releasekitkat-mr2.2-releasekitkat-mr2.1-releasekitkat-mr2-releasekitkat-mr1.1-releasekitkat-mr1-releasekitkat-devkitkat-cts-releasekitkat-cts-dev

wait for AccessibilityEvent after injecting events
via UiAutomation

Change-Id: I3d56e07cf2e7912a21de12d1a7bacd4f33e1bc5a

18 files changed, 350 insertions, 88 deletions

diff --git a/src/com/google/android/droiddriver/DroidDriver.java b/src/com/google/android/droiddriver/DroidDriver.java
index 03aa4e6..6daed71 100644
--- a/src/com/google/android/droiddriver/DroidDriver.java
+++ b/src/com/google/android/droiddriver/DroidDriver.java
@@ -106,7 +106,10 @@ public interface DroidDriver {
   void setPoller(Poller poller);
 
   /**
-   * Dumps the UiElement tree to a file to help debug.
+   * Dumps the UiElement tree to a file to help debug. The tree is based on the
+   * last used root UiElement if it exists. Screenshot is always current. If
+   * they do not match, the UiElement tree must be stale, indicating that you
+   * should use a fresh UiElement instead of an old instance.
    *
    * @param path the path of file to save the tree
    * @return whether the dumping succeeded
diff --git a/src/com/google/android/droiddriver/InputInjector.java b/src/com/google/android/droiddriver/InputInjector.java
index 3b2d675..987bb7d 100644
--- a/src/com/google/android/droiddriver/InputInjector.java
+++ b/src/com/google/android/droiddriver/InputInjector.java
@@ -19,12 +19,12 @@ package com.google.android.droiddriver;
 import android.view.InputEvent;
 
 /**
- * Interface for interacting with the UI at via InputEventInjection.
+ * Interface for interacting with the UI via InputEvent injection.
  */
 public interface InputInjector {
 
   /**
-   * Injects the InputEvent.
+   * Injects the {@code event}.
    *
    * @param event The event to inject.
    * @return true if the injection succeeded.
diff --git a/src/com/google/android/droiddriver/actions/Action.java b/src/com/google/android/droiddriver/actions/Action.java
index 80ec37d..1c57e12 100644
--- a/src/com/google/android/droiddriver/actions/Action.java
+++ b/src/com/google/android/droiddriver/actions/Action.java
@@ -16,17 +16,20 @@
 
 package com.google.android.droiddriver.actions;
 
+import android.view.InputEvent;
+
 import com.google.android.droiddriver.InputInjector;
 import com.google.android.droiddriver.UiElement;
 
 /**
- * Interface for performing action on a UiElement.
+ * Interface for performing action on a UiElement. An action is a high-level
+ * user interaction that consists of a series of {@link InputEvent}s.
  */
 public interface Action {
   /**
    * Performs the action.
    *
-   * @param injector the injector to inject input events
+   * @param injector the injector to inject {@link InputEvent}s
    * @param element the Ui element to perform the action on
    * @return Whether the action is successful. Some actions throw exceptions in
    *         case of failure, when that behavior is more appropriate. For
@@ -35,6 +38,15 @@ public interface Action {
   boolean perform(InputInjector injector, UiElement element);
 
   /**
+   * Gets the timeout to wait for an indicator that the action has been carried
+   * out. Different DroidDriver implementations use this value in different
+   * ways. For example, UiAutomationDriver waits for AccessibilityEvent up to
+   * this value. InstrumentationDriver ignores this value because it
+   * synchronizes on the event loop.
+   */
+  long getTimeoutMillis();
+
+  /**
    * {@inheritDoc}
    *
    * <p>
diff --git a/src/com/google/android/droiddriver/actions/BaseAction.java b/src/com/google/android/droiddriver/actions/BaseAction.java
new file mode 100644
index 0000000..a469fec
--- /dev/null
+++ b/src/com/google/android/droiddriver/actions/BaseAction.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2013 DroidDriver committers
+ *
+ * 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.google.android.droiddriver.actions;
+
+/**
+ * Base class of {@link Action} that implements {@link #getTimeoutMillis}.
+ */
+public abstract class BaseAction implements Action {
+  private final long timeoutMillis;
+
+  @Override
+  public long getTimeoutMillis() {
+    return timeoutMillis;
+  }
+
+  protected BaseAction(long timeoutMillis) {
+    this.timeoutMillis = timeoutMillis;
+  }
+}
diff --git a/src/com/google/android/droiddriver/actions/ClickAction.java b/src/com/google/android/droiddriver/actions/ClickAction.java
index ea954e1..12db023 100644
--- a/src/com/google/android/droiddriver/actions/ClickAction.java
+++ b/src/com/google/android/droiddriver/actions/ClickAction.java
@@ -22,28 +22,37 @@ import android.view.ViewConfiguration;
 
 import com.google.android.droiddriver.InputInjector;
 import com.google.android.droiddriver.UiElement;
-import com.google.android.droiddriver.exceptions.ActionException;
 import com.google.android.droiddriver.util.Events;
 
 /**
  * An action that does clicks on an UiElement.
  */
-public enum ClickAction implements Action {
-  SINGLE {
-    /** @throws ActionException */
+public abstract class ClickAction extends BaseAction {
+
+  public static final ClickAction SINGLE = new SingleClick(1000L);
+  public static final ClickAction LONG = new LongClick(1000L);
+  public static final ClickAction DOUBLE = new DoubleClick(1000L);
+
+  private static final long CLICK_DURATION_MILLIS = 100L;
+
+  public static class DoubleClick extends ClickAction {
+    public DoubleClick(long timeoutMillis) {
+      super(timeoutMillis);
+    }
+
     @Override
     public boolean perform(InputInjector injector, UiElement element) {
-      Rect elementRect = element.getVisibleBounds();
-      long downTime = Events.touchDown(injector, elementRect.centerX(), elementRect.centerY());
-      // UiAutomator clickAndSync does this, while
-      // android.test.TouchUtils#clickView sleep 1000
-      SystemClock.sleep(CLICK_DURATION_MILLIS);
-      Events.touchUp(injector, downTime, elementRect.centerX(), elementRect.centerY());
+      SINGLE.perform(injector, element);
+      SINGLE.perform(injector, element);
       return true;
     }
-  },
-  LONG {
-    /** @throws ActionException */
+  }
+
+  private static class LongClick extends ClickAction {
+    public LongClick(long timeoutMillis) {
+      super(timeoutMillis);
+    }
+
     @Override
     public boolean perform(InputInjector injector, UiElement element) {
       Rect elementRect = element.getVisibleBounds();
@@ -53,16 +62,31 @@ public enum ClickAction implements Action {
       Events.touchUp(injector, downTime, elementRect.centerX(), elementRect.centerY());
       return true;
     }
-  },
-  DOUBLE {
-    /** @throws ActionException */
+  }
+
+  public static class SingleClick extends ClickAction {
+    public SingleClick(long timeoutMillis) {
+      super(timeoutMillis);
+    }
+
     @Override
     public boolean perform(InputInjector injector, UiElement element) {
-      SINGLE.perform(injector, element);
-      SINGLE.perform(injector, element);
+      Rect elementRect = element.getVisibleBounds();
+      long downTime = Events.touchDown(injector, elementRect.centerX(), elementRect.centerY());
+      // UiAutomator clickAndSync does this, while
+      // android.test.TouchUtils#clickView sleep 1000
+      SystemClock.sleep(CLICK_DURATION_MILLIS);
+      Events.touchUp(injector, downTime, elementRect.centerX(), elementRect.centerY());
       return true;
     }
-  };
+  }
+
+  protected ClickAction(long timeoutMillis) {
+    super(timeoutMillis);
+  }
 
-  private static final long CLICK_DURATION_MILLIS = 100;
+  @Override
+  public String toString() {
+    return getClass().getSimpleName();
+  }
 }
diff --git a/src/com/google/android/droiddriver/actions/KeyAction.java b/src/com/google/android/droiddriver/actions/KeyAction.java
new file mode 100644
index 0000000..1cafe0e
--- /dev/null
+++ b/src/com/google/android/droiddriver/actions/KeyAction.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2013 DroidDriver committers
+ *
+ * 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.google.android.droiddriver.actions;
+
+/**
+ * Base class for {@link Action} that injects key events.
+ */
+public abstract class KeyAction extends BaseAction {
+  protected KeyAction(long timeoutMillis) {
+    super(timeoutMillis);
+  }
+}
diff --git a/src/com/google/android/droiddriver/actions/PressKeyAction.java b/src/com/google/android/droiddriver/actions/PressKeyAction.java
index 6cc901b..14fd060 100644
--- a/src/com/google/android/droiddriver/actions/PressKeyAction.java
+++ b/src/com/google/android/droiddriver/actions/PressKeyAction.java
@@ -25,12 +25,20 @@ import com.google.android.droiddriver.util.Events;
 import com.google.common.base.Objects;
 
 /**
- * An general action to press any of the soft keys on the device.
+ * An action to press a single key. TODO: rename to SingleKeyAction
  */
-public class PressKeyAction implements Action {
+public class PressKeyAction extends KeyAction {
   private final int keyCode;
 
+  /**
+   * Defaults timeoutMillis to 0.
+   */
   public PressKeyAction(int keyCode) {
+    this(keyCode, 0L);
+  }
+
+  public PressKeyAction(int keyCode, long timeoutMillis) {
+    super(timeoutMillis);
     this.keyCode = keyCode;
   }
 
diff --git a/src/com/google/android/droiddriver/actions/ScrollAction.java b/src/com/google/android/droiddriver/actions/ScrollAction.java
new file mode 100644
index 0000000..e616e09
--- /dev/null
+++ b/src/com/google/android/droiddriver/actions/ScrollAction.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2013 DroidDriver committers
+ *
+ * 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.google.android.droiddriver.actions;
+
+/**
+ * Base class for {@link Action} that scrolls.
+ */
+public abstract class ScrollAction extends BaseAction {
+  protected ScrollAction(long timeoutMillis) {
+    super(timeoutMillis);
+  }
+}
diff --git a/src/com/google/android/droiddriver/actions/SwipeAction.java b/src/com/google/android/droiddriver/actions/SwipeAction.java
index 116b828..433f6de 100644
--- a/src/com/google/android/droiddriver/actions/SwipeAction.java
+++ b/src/com/google/android/droiddriver/actions/SwipeAction.java
@@ -26,14 +26,22 @@ import com.google.android.droiddriver.exceptions.ActionException;
 import com.google.android.droiddriver.util.Events;
 
 /**
- * An action that does a swipe.
+ * A {@link ScrollAction} that swipes the touch screen.
  */
-public class SwipeAction implements Action {
+public class SwipeAction extends ScrollAction {
 
   private final ScrollDirection direction;
   private final boolean drag;
 
+  /**
+   * Defaults timeoutMillis to 0.
+   */
   public SwipeAction(ScrollDirection direction, boolean drag) {
+    this(direction, drag, 0L);
+  }
+
+  public SwipeAction(ScrollDirection direction, boolean drag, long timeoutMillis) {
+    super(timeoutMillis);
     this.direction = direction;
     this.drag = drag;
   }
diff --git a/src/com/google/android/droiddriver/actions/TypeAction.java b/src/com/google/android/droiddriver/actions/TypeAction.java
index 2fb182b..55e8dea 100644
--- a/src/com/google/android/droiddriver/actions/TypeAction.java
+++ b/src/com/google/android/droiddriver/actions/TypeAction.java
@@ -29,14 +29,22 @@ import android.view.KeyEvent;
 /**
  * An action to type text.
  */
-public class TypeAction implements Action {
+public class TypeAction extends KeyAction {
 
   private static final KeyCharacterMap KEY_CHAR_MAP = KeyCharacterMap
       .load(KeyCharacterMap.VIRTUAL_KEYBOARD);
 
   private final String text;
 
+  /**
+   * Defaults timeoutMillis to 0.
+   */
   public TypeAction(String text) {
+    this(text, 0L);
+  }
+
+  public TypeAction(String text, long timeoutMillis) {
+    super(timeoutMillis);
     this.text = Preconditions.checkNotNull(text);
   }
 
diff --git a/src/com/google/android/droiddriver/base/AbstractContext.java b/src/com/google/android/droiddriver/base/AbstractContext.java
new file mode 100644
index 0000000..2c831ee
--- /dev/null
+++ b/src/com/google/android/droiddriver/base/AbstractContext.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2013 DroidDriver committers
+ *
+ * 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.google.android.droiddriver.base;
+
+import com.google.android.droiddriver.InputInjector;
+
+/**
+ * Internal helper for managing all instances.
+ */
+public abstract class AbstractContext {
+  protected final InputInjector injector;
+
+  protected AbstractContext(InputInjector injector) {
+    this.injector = injector;
+  }
+
+  public InputInjector getInjector() {
+    return injector;
+  }
+
+  /** Clears data in the context */
+  public abstract void clearData();
+}
diff --git a/src/com/google/android/droiddriver/base/AbstractDroidDriver.java b/src/com/google/android/droiddriver/base/AbstractDroidDriver.java
index 63b355c..ed78b46 100644
--- a/src/com/google/android/droiddriver/base/AbstractDroidDriver.java
+++ b/src/com/google/android/droiddriver/base/AbstractDroidDriver.java
@@ -16,6 +16,7 @@
 
 package com.google.android.droiddriver.base;
 
+import android.app.Instrumentation;
 import android.graphics.Bitmap;
 import android.graphics.Bitmap.CompressFormat;
 import android.util.Log;
@@ -31,6 +32,7 @@ import com.google.android.droiddriver.finders.Finder;
 import com.google.android.droiddriver.util.DefaultPoller;
 import com.google.android.droiddriver.util.FileUtils;
 import com.google.android.droiddriver.util.Logs;
+import com.google.common.base.Preconditions;
 
 import java.io.BufferedOutputStream;
 
@@ -40,12 +42,18 @@ import java.io.BufferedOutputStream;
  */
 public abstract class AbstractDroidDriver implements DroidDriver, Screenshotter {
 
+  protected final Instrumentation instrumentation;
   private Poller poller = new DefaultPoller();
+  private AbstractUiElement rootElement;
+
+  protected AbstractDroidDriver(Instrumentation instrumentation) {
+    this.instrumentation = Preconditions.checkNotNull(instrumentation);
+  }
 
   @Override
   public UiElement find(Finder finder) {
     Logs.call(this, "find", finder);
-    return finder.find(getRootElement());
+    return finder.find(refreshRootElement());
   }
 
   @Override
@@ -96,7 +104,22 @@ public abstract class AbstractDroidDriver implements DroidDriver, Screenshotter
     this.poller = poller;
   }
 
-  protected abstract AbstractUiElement getRootElement();
+  protected abstract AbstractUiElement getNewRootElement();
+
+  protected abstract AbstractContext getContext();
+
+  protected AbstractUiElement getRootElement() {
+    if (rootElement == null) {
+      refreshRootElement();
+    }
+    return rootElement;
+  }
+
+  private AbstractUiElement refreshRootElement() {
+    getContext().clearData();
+    rootElement = getNewRootElement();
+    return rootElement;
+  }
 
   @Override
   public boolean dumpUiElementTree(String path) {
diff --git a/src/com/google/android/droiddriver/base/AbstractUiElement.java b/src/com/google/android/droiddriver/base/AbstractUiElement.java
index c82a360..65048f2 100644
--- a/src/com/google/android/droiddriver/base/AbstractUiElement.java
+++ b/src/com/google/android/droiddriver/base/AbstractUiElement.java
@@ -37,6 +37,8 @@ import org.w3c.dom.Element;
 
 import java.lang.ref.WeakReference;
 import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.FutureTask;
 
 /**
  * Abstract implementation with common methods already implemented.
@@ -53,9 +55,39 @@ public abstract class AbstractUiElement implements UiElement {
   public boolean perform(Action action) {
     Logs.call(this, "perform", action);
     checkVisible();
+    return performAndWait(action);
+  }
+
+  protected boolean doPerform(Action action) {
     return action.perform(getInjector(), this);
   }
 
+  protected void doPerformAndWait(FutureTask<Boolean> futureTask, long timeoutMillis) {
+    // ignores timeoutMillis; subclasses can override this behavior
+    futureTask.run();
+  }
+
+  private boolean performAndWait(final Action action) {
+    // timeoutMillis <= 0 means no need to wait
+    if (action.getTimeoutMillis() <= 0) {
+      return doPerform(action);
+    }
+
+    FutureTask<Boolean> futureTask = new FutureTask<Boolean>(new Callable<Boolean>() {
+      @Override
+      public Boolean call() {
+        return doPerform(action);
+      }
+    });
+    doPerformAndWait(futureTask, action.getTimeoutMillis());
+    try {
+      return futureTask.get();
+    } catch (Exception e) {
+      // should not reach here b/c futureTask has run
+      return false;
+    }
+  }
+
   @Override
   public void setText(String text) {
     // TODO: Define common actions as a const.
diff --git a/src/com/google/android/droiddriver/instrumentation/InstrumentationContext.java b/src/com/google/android/droiddriver/instrumentation/InstrumentationContext.java
index fed7ee6..0c4e9dc 100644
--- a/src/com/google/android/droiddriver/instrumentation/InstrumentationContext.java
+++ b/src/com/google/android/droiddriver/instrumentation/InstrumentationContext.java
@@ -23,8 +23,8 @@ import android.view.MotionEvent;
 import android.view.View;
 
 import com.google.android.droiddriver.InputInjector;
+import com.google.android.droiddriver.base.AbstractContext;
 import com.google.android.droiddriver.exceptions.ActionException;
-import com.google.common.base.Preconditions;
 import com.google.common.collect.MapMaker;
 
 import java.util.Map;
@@ -32,34 +32,23 @@ import java.util.Map;
 /**
  * Internal helper for managing all instances.
  */
-public class InstrumentationContext {
-  private final Instrumentation instrumentation;
-  private final InputInjector injector;
+public class InstrumentationContext extends AbstractContext {
   private final Map<View, ViewElement> map = new MapMaker().weakKeys().weakValues().makeMap();
 
-  InstrumentationContext(Instrumentation instrumentation) {
-    this.instrumentation = Preconditions.checkNotNull(instrumentation);
-    injector = new InputInjector() {
+  InstrumentationContext(final Instrumentation instrumentation) {
+    super(new InputInjector() {
       @Override
       public boolean injectInputEvent(InputEvent event) {
         if (event instanceof MotionEvent) {
-          getInstrumentation().sendPointerSync((MotionEvent) event);
+          instrumentation.sendPointerSync((MotionEvent) event);
         } else if (event instanceof KeyEvent) {
-          getInstrumentation().sendKeySync((KeyEvent) event);
+          instrumentation.sendKeySync((KeyEvent) event);
         } else {
           throw new ActionException("Unknown input event type: " + event);
         }
         return true;
       }
-    };
-  }
-
-  public Instrumentation getInstrumentation() {
-    return instrumentation;
-  }
-
-  public InputInjector getInjector() {
-    return injector;
+    });
   }
 
   public ViewElement getUiElement(View view) {
@@ -70,4 +59,9 @@ public class InstrumentationContext {
     }
     return element;
   }
+
+  @Override
+  public void clearData() {
+    map.clear();
+  }
 }
diff --git a/src/com/google/android/droiddriver/instrumentation/InstrumentationDriver.java b/src/com/google/android/droiddriver/instrumentation/InstrumentationDriver.java
index 800f884..c4f1fbb 100644
--- a/src/com/google/android/droiddriver/instrumentation/InstrumentationDriver.java
+++ b/src/com/google/android/droiddriver/instrumentation/InstrumentationDriver.java
@@ -31,18 +31,23 @@ import com.google.common.primitives.Longs;
  * Implementation of a UiDriver that is driven via instrumentation.
  */
 public class InstrumentationDriver extends AbstractDroidDriver {
-
   private final InstrumentationContext context;
 
   public InstrumentationDriver(Instrumentation instrumentation) {
+    super(instrumentation);
     this.context = new InstrumentationContext(instrumentation);
   }
 
   @Override
-  public ViewElement getRootElement() {
+  protected ViewElement getNewRootElement() {
     return context.getUiElement(findRootView());
   }
 
+  @Override
+  protected InstrumentationContext getContext() {
+    return context;
+  }
+
   private View findRootView() {
     Activity runningActivity = getRunningActivity();
     View[] views = RootFinder.getRootViews();
@@ -60,7 +65,7 @@ public class InstrumentationDriver extends AbstractDroidDriver {
     long timeoutMillis = getPoller().getTimeoutMillis();
     long end = SystemClock.uptimeMillis() + timeoutMillis;
     while (true) {
-      context.getInstrumentation().waitForIdleSync();
+      instrumentation.waitForIdleSync();
       Activity runningActivity = ActivityUtils.getRunningActivity();
       if (runningActivity != null) {
         return runningActivity;
@@ -94,7 +99,7 @@ public class InstrumentationDriver extends AbstractDroidDriver {
   @Override
   protected Bitmap takeScreenshot() {
     ScreenshotRunnable screenshotRunnable = new ScreenshotRunnable(findRootView());
-    context.getInstrumentation().runOnMainSync(screenshotRunnable);
+    instrumentation.runOnMainSync(screenshotRunnable);
     return screenshotRunnable.screenshot;
   }
 }
diff --git a/src/com/google/android/droiddriver/uiautomation/UiAutomationContext.java b/src/com/google/android/droiddriver/uiautomation/UiAutomationContext.java
index cfd1488..7efa48e 100644
--- a/src/com/google/android/droiddriver/uiautomation/UiAutomationContext.java
+++ b/src/com/google/android/droiddriver/uiautomation/UiAutomationContext.java
@@ -16,13 +16,12 @@
 
 package com.google.android.droiddriver.uiautomation;
 
-import android.app.Instrumentation;
 import android.app.UiAutomation;
 import android.view.InputEvent;
 import android.view.accessibility.AccessibilityNodeInfo;
 
 import com.google.android.droiddriver.InputInjector;
-import com.google.common.base.Preconditions;
+import com.google.android.droiddriver.base.AbstractContext;
 import com.google.common.collect.MapMaker;
 
 import java.util.Map;
@@ -30,33 +29,20 @@ import java.util.Map;
 /**
  * Internal helper for managing all instances.
  */
-public class UiAutomationContext {
-  private final Instrumentation instrumentation;
-  private final InputInjector injector;
+public class UiAutomationContext extends AbstractContext {
   // Maybe we should use Cache instead of Map on memory-constrained devices
   private final Map<AccessibilityNodeInfo, UiAutomationElement> map = new MapMaker().weakKeys()
       .weakValues().makeMap();
+  private final UiAutomation uiAutomation;
 
-  UiAutomationContext(Instrumentation instrumentation) {
-    this.instrumentation = Preconditions.checkNotNull(instrumentation);
-    injector = new InputInjector() {
+  UiAutomationContext(final UiAutomation uiAutomation) {
+    super(new InputInjector() {
       @Override
       public boolean injectInputEvent(InputEvent event) {
-        return getUiAutomation().injectInputEvent(event, true /* sync */);
+        return uiAutomation.injectInputEvent(event, true /* sync */);
       }
-    };
-  }
-
-  public Instrumentation getInstrumentation() {
-    return instrumentation;
-  }
-
-  public UiAutomation getUiAutomation() {
-    return instrumentation.getUiAutomation();
-  }
-
-  public InputInjector getInjector() {
-    return injector;
+    });
+    this.uiAutomation = uiAutomation;
   }
 
   public UiAutomationElement getUiElement(AccessibilityNodeInfo node) {
@@ -67,4 +53,13 @@ public class UiAutomationContext {
     }
     return element;
   }
+
+  @Override
+  public void clearData() {
+    map.clear();
+  }
+
+  public UiAutomation getUiAutomation() {
+    return uiAutomation;
+  }
 }
diff --git a/src/com/google/android/droiddriver/uiautomation/UiAutomationDriver.java b/src/com/google/android/droiddriver/uiautomation/UiAutomationDriver.java
index 03c8969..7d21233 100644
--- a/src/com/google/android/droiddriver/uiautomation/UiAutomationDriver.java
+++ b/src/com/google/android/droiddriver/uiautomation/UiAutomationDriver.java
@@ -39,29 +39,34 @@ public class UiAutomationDriver extends AbstractDroidDriver {
   private static final long QUIET_TIME_TO_BE_CONSIDERD_IDLE_STATE = 500;// ms
 
   private final UiAutomationContext context;
-  private final Instrumentation instrumentation;
+  private final UiAutomation uiAutomation;
 
   public UiAutomationDriver(Instrumentation instrumentation) {
-    this.instrumentation = instrumentation;
-    this.context = new UiAutomationContext(instrumentation);
+    super(instrumentation);
+    this.uiAutomation = instrumentation.getUiAutomation();
+    this.context = new UiAutomationContext(uiAutomation);
   }
 
   @Override
-  public UiAutomationElement getRootElement() {
+  protected UiAutomationElement getNewRootElement() {
     return context.getUiElement(getRootNode());
   }
 
+  @Override
+  protected UiAutomationContext getContext() {
+    return context;
+  }
+
   private AccessibilityNodeInfo getRootNode() {
     long timeoutMillis = getPoller().getTimeoutMillis();
     try {
-      getUiAutomation().waitForIdle(QUIET_TIME_TO_BE_CONSIDERD_IDLE_STATE,
-              timeoutMillis);
+      uiAutomation.waitForIdle(QUIET_TIME_TO_BE_CONSIDERD_IDLE_STATE, timeoutMillis);
     } catch (java.util.concurrent.TimeoutException e) {
       throw new TimeoutException(e);
     }
     long end = SystemClock.uptimeMillis() + timeoutMillis;
     while (true) {
-      AccessibilityNodeInfo root = getUiAutomation().getRootInActiveWindow();
+      AccessibilityNodeInfo root = uiAutomation.getRootInActiveWindow();
       if (root != null) {
         return root;
       }
@@ -77,10 +82,6 @@ public class UiAutomationDriver extends AbstractDroidDriver {
 
   @Override
   protected Bitmap takeScreenshot() {
-    return getUiAutomation().takeScreenshot();
-  }
-
-  private UiAutomation getUiAutomation() {
-      return instrumentation.getUiAutomation();
+    return uiAutomation.takeScreenshot();
   }
 }
diff --git a/src/com/google/android/droiddriver/uiautomation/UiAutomationElement.java b/src/com/google/android/droiddriver/uiautomation/UiAutomationElement.java
index 9773452..a399360 100644
--- a/src/com/google/android/droiddriver/uiautomation/UiAutomationElement.java
+++ b/src/com/google/android/droiddriver/uiautomation/UiAutomationElement.java
@@ -18,8 +18,11 @@ package com.google.android.droiddriver.uiautomation;
 
 import static com.google.android.droiddriver.util.TextUtils.charSequenceToString;
 
+import android.app.UiAutomation;
+import android.app.UiAutomation.AccessibilityEventFilter;
 import android.graphics.Rect;
 import android.util.Log;
+import android.view.accessibility.AccessibilityEvent;
 import android.view.accessibility.AccessibilityNodeInfo;
 
 import com.google.android.droiddriver.InputInjector;
@@ -27,16 +30,28 @@ import com.google.android.droiddriver.base.AbstractUiElement;
 import com.google.android.droiddriver.util.Logs;
 import com.google.common.base.Preconditions;
 
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeoutException;
+
 /**
  * A UiElement that is backed by the UiAutomation object.
  */
 public class UiAutomationElement extends AbstractUiElement {
+  private static final AccessibilityEventFilter ANY_EVENT_FILTER = new AccessibilityEventFilter() {
+    @Override
+    public boolean accept(AccessibilityEvent arg0) {
+      return true;
+    }
+  };
+
   private final UiAutomationContext context;
   private final AccessibilityNodeInfo node;
+  private final UiAutomation uiAutomation;
 
   public UiAutomationElement(UiAutomationContext context, AccessibilityNodeInfo node) {
     this.context = Preconditions.checkNotNull(context);
     this.node = Preconditions.checkNotNull(node);
+    this.uiAutomation = context.getUiAutomation();
   }
 
   @Override
@@ -164,4 +179,16 @@ public class UiAutomationElement extends AbstractUiElement {
     AccessibilityNodeInfo parent = node.getParent();
     return parent == null ? null : context.getUiElement(parent);
   }
+
+  @Override
+  protected void doPerformAndWait(FutureTask<Boolean> futureTask, long timeoutMillis) {
+    try {
+      uiAutomation.executeAndWaitForEvent(futureTask, ANY_EVENT_FILTER, timeoutMillis);
+    } catch (TimeoutException e) {
+      // This is for sync'ing with Accessibility API on best-effort because
+      // it is not reliable.
+      // Exception is ignored here. Tests will fail anyways if this is
+      // critical.
+    }
+  }
 }