Merge "Fix AssistDataRequesterTest"
diff --git a/services/core/java/com/android/server/am/AssistDataRequester.java b/services/core/java/com/android/server/am/AssistDataRequester.java
index 09df7e20..d8d6528 100644
--- a/services/core/java/com/android/server/am/AssistDataRequester.java
+++ b/services/core/java/com/android/server/am/AssistDataRequester.java
@@ -24,6 +24,7 @@
 
 import android.app.ActivityTaskManager;
 import android.app.AppOpsManager;
+import android.app.IActivityTaskManager;
 import android.app.IAssistDataReceiver;
 import android.content.Context;
 import android.graphics.Bitmap;
@@ -33,6 +34,7 @@
 import android.view.IWindowManager;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.MetricsLogger;
 
 import java.io.PrintWriter;
@@ -50,6 +52,8 @@
     public static final String KEY_RECEIVER_EXTRA_INDEX = "index";
 
     private IWindowManager mWindowManager;
+    @VisibleForTesting
+    public IActivityTaskManager mActivityTaskManager;
     private Context mContext;
     private AppOpsManager mAppOpsManager;
 
@@ -128,6 +132,7 @@
         mCallbacks = callbacks;
         mCallbacksLock = callbacksLock;
         mWindowManager = windowManager;
+        mActivityTaskManager = ActivityTaskManager.getService();
         mContext = context;
         mAppOpsManager = appOpsManager;
         mRequestStructureAppOps = requestStructureAppOps;
@@ -195,8 +200,7 @@
         // Ensure that the current activity supports assist data
         boolean isAssistDataAllowed = false;
         try {
-            isAssistDataAllowed =
-                    ActivityTaskManager.getService().isAssistDataAllowedOnCurrentActivity();
+            isAssistDataAllowed = mActivityTaskManager.isAssistDataAllowedOnCurrentActivity();
         } catch (RemoteException e) {
             // Should never happen
         }
@@ -222,9 +226,9 @@
                         receiverExtras.putInt(KEY_RECEIVER_EXTRA_INDEX, i);
                         receiverExtras.putInt(KEY_RECEIVER_EXTRA_COUNT, numActivities);
                         boolean result = requestAutofillData
-                                ? ActivityTaskManager.getService().requestAutofillData(this,
+                                ? mActivityTaskManager.requestAutofillData(this,
                                         receiverExtras, topActivity, 0 /* flags */)
-                                : ActivityTaskManager.getService().requestAssistContextExtras(
+                                : mActivityTaskManager.requestAssistContextExtras(
                                         ASSIST_CONTEXT_FULL, this, receiverExtras, topActivity,
                                         /* focused= */ i == 0, /* newSessionId= */ i == 0);
                         if (result) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/AssistDataRequesterTest.java b/services/tests/wmtests/src/com/android/server/wm/AssistDataRequesterTest.java
index 1e02a12..329af95 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AssistDataRequesterTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AssistDataRequesterTest.java
@@ -38,9 +38,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 
 import android.app.AppOpsManager;
-import android.app.IActivityManager;
 import android.app.IActivityTaskManager;
-import android.content.Context;
 import android.graphics.Bitmap;
 import android.os.Bundle;
 import android.os.Handler;
@@ -50,7 +48,6 @@
 import android.util.Log;
 import android.view.IWindowManager;
 
-import androidx.test.filters.FlakyTest;
 import androidx.test.filters.MediumTest;
 
 import com.android.server.am.AssistDataRequester;
@@ -88,12 +85,10 @@
     private static final int TEST_UID = 0;
     private static final String TEST_PACKAGE = "";
 
-    private Context mContext;
     private AssistDataRequester mDataRequester;
     private Callbacks mCallbacks;
     private Object mCallbacksLock;
     private Handler mHandler;
-    private IActivityManager mAm;
     private IActivityTaskManager mAtm;
     private IWindowManager mWm;
     private AppOpsManager mAppOpsManager;
@@ -109,7 +104,6 @@
 
     @Before
     public void setUp() throws Exception {
-        mAm = mock(IActivityManager.class);
         mAtm = mock(IActivityTaskManager.class);
         mWm = mock(IWindowManager.class);
         mAppOpsManager = mock(AppOpsManager.class);
@@ -118,7 +112,7 @@
         mCallbacks = new Callbacks();
         mDataRequester = new AssistDataRequester(mContext, mWm, mAppOpsManager, mCallbacks,
                 mCallbacksLock, OP_ASSIST_STRUCTURE, OP_ASSIST_SCREENSHOT);
-
+        mDataRequester.mActivityTaskManager = mAtm;
         // Gate the continuation of the assist data callbacks until we are ready within the tests
         mGate = new CountDownLatch(1);
         doAnswer(invocation -> {
@@ -155,7 +149,6 @@
                 .checkOpNoThrow(eq(OP_ASSIST_SCREENSHOT), anyInt(), anyString());
     }
 
-    @FlakyTest(bugId = 124088319)
     @Test
     public void testRequestData() throws Exception {
         setupMocks(CURRENT_ACTIVITY_ASSIST_ALLOWED, CALLER_ASSIST_STRUCTURE_ALLOWED,
@@ -176,7 +169,6 @@
         assertReceivedDataCount(0, 0, 0, 0);
     }
 
-    @FlakyTest(bugId = 124088319)
     @Test
     public void testCurrentAppDisallow_expectNullCallbacks() throws Exception {
         setupMocks(!CURRENT_ACTIVITY_ASSIST_ALLOWED, CALLER_ASSIST_STRUCTURE_ALLOWED,
@@ -187,7 +179,6 @@
         assertReceivedDataCount(0, 1, 0, 1);
     }
 
-    @FlakyTest(bugId = 124088319)
     @Test
     public void testProcessPendingData() throws Exception {
         setupMocks(CURRENT_ACTIVITY_ASSIST_ALLOWED, CALLER_ASSIST_STRUCTURE_ALLOWED,
@@ -245,7 +236,6 @@
         assertReceivedDataCount(0, 1, 0, 1);
     }
 
-    @FlakyTest(bugId = 124088319)
     @Test
     public void testDisallowAssistContextExtras_expectNullDataCallbacks() throws Exception {
         setupMocks(CURRENT_ACTIVITY_ASSIST_ALLOWED, CALLER_ASSIST_STRUCTURE_ALLOWED,
@@ -259,7 +249,6 @@
         assertReceivedDataCount(0, 1, 0, 1);
     }
 
-    @FlakyTest(bugId = 124088319)
     @Test
     public void testNoFetchScreenshots_expectNoScreenshotCallbacks() throws Exception {
         setupMocks(CURRENT_ACTIVITY_ASSIST_ALLOWED, CALLER_ASSIST_STRUCTURE_ALLOWED,
@@ -270,7 +259,6 @@
         assertReceivedDataCount(5, 5, 0, 0);
     }
 
-    @FlakyTest(bugId = 124088319)
     @Test
     public void testDisallowAssistScreenshot_expectNullScreenshotCallback() throws Exception {
         setupMocks(CURRENT_ACTIVITY_ASSIST_ALLOWED, CALLER_ASSIST_STRUCTURE_ALLOWED,