fix memory leak in Contexts UiElement cache

UiElement values need to be WeakReference as well

Change-Id: I31f6e833216154a3d3da7ebabc856d8ce8221631

1 files changed, 3 insertions, 2 deletions

diff --git a/src/com/google/android/droiddriver/uiautomation/UiAutomationContext.java b/src/com/google/android/droiddriver/uiautomation/UiAutomationContext.java
index b171ac6..cfd1488 100644
--- a/src/com/google/android/droiddriver/uiautomation/UiAutomationContext.java
+++ b/src/com/google/android/droiddriver/uiautomation/UiAutomationContext.java
@@ -33,8 +33,9 @@ import java.util.Map;
 public class UiAutomationContext {
   private final Instrumentation instrumentation;
   private final InputInjector injector;
+  // Maybe we should use Cache instead of Map on memory-constrained devices
   private final Map<AccessibilityNodeInfo, UiAutomationElement> map = new MapMaker().weakKeys()
-      .makeMap();
+      .weakValues().makeMap();
 
   UiAutomationContext(Instrumentation instrumentation) {
     this.instrumentation = Preconditions.checkNotNull(instrumentation);
@@ -51,7 +52,7 @@ public class UiAutomationContext {
   }
 
   public UiAutomation getUiAutomation() {
-      return instrumentation.getUiAutomation();
+    return instrumentation.getUiAutomation();
   }
 
   public InputInjector getInjector() {