Merge SP2A.220305.013

Bug: 220074017
Change-Id: Idfdd94e902f656ac65a2a75dfdd199f6f85ba472

34 files changed, 2141 insertions, 247 deletions

diff --git a/core/tests/BroadcastRadioTests/OWNERS b/core/tests/BroadcastRadioTests/OWNERS
index ea4421eae96a..3e360e7e992c 100644
--- a/core/tests/BroadcastRadioTests/OWNERS
+++ b/core/tests/BroadcastRadioTests/OWNERS
@@ -1,2 +1,3 @@
+keunyoung@google.com
+oscarazu@google.com
 twasilczyk@google.com
-randolphs@google.com
diff --git a/core/tests/BroadcastRadioTests/src/com/android/server/broadcastradio/hal2/StartProgramListUpdatesFanoutTest.java b/core/tests/BroadcastRadioTests/src/com/android/server/broadcastradio/hal2/StartProgramListUpdatesFanoutTest.java
index 7c6271cbdf61..c194989b2752 100644
--- a/core/tests/BroadcastRadioTests/src/com/android/server/broadcastradio/hal2/StartProgramListUpdatesFanoutTest.java
+++ b/core/tests/BroadcastRadioTests/src/com/android/server/broadcastradio/hal2/StartProgramListUpdatesFanoutTest.java
@@ -65,6 +65,7 @@ public class StartProgramListUpdatesFanoutTest {
     @Mock ITunerSession mHalTunerSessionMock;
     private android.hardware.radio.ITunerCallback[] mAidlTunerCallbackMocks;
 
+    private final Object mLock = new Object();
     // RadioModule under test
     private RadioModule mRadioModule;
 
@@ -96,7 +97,7 @@ public class StartProgramListUpdatesFanoutTest {
 
         mRadioModule = new RadioModule(mBroadcastRadioMock, new RadioManager.ModuleProperties(0, "",
                   0, "", "", "", "", 0, 0, false, false, null, false, new int[] {}, new int[] {},
-                  null, null));
+                  null, null), mLock);
 
         doAnswer((Answer) invocation -> {
             mHalTunerCallback = (ITunerCallback) invocation.getArguments()[0];
diff --git a/core/tests/coretests/res/drawable/custom_drawable.xml b/core/tests/coretests/res/drawable/custom_drawable.xml
new file mode 100644
index 000000000000..ebb821fa11fb
--- /dev/null
+++ b/core/tests/coretests/res/drawable/custom_drawable.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Copyright (C) 2021 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.
+  -->
+
+<drawable xmlns:android="http://schemas.android.com/apk/res/android"
+          class="android.window.CustomDrawable"
+          android:drawable="@drawable/bitmap_drawable"
+          android:inset="10dp" />
\ No newline at end of file
diff --git a/core/tests/coretests/src/android/accessibilityservice/AccessibilityServiceTest.java b/core/tests/coretests/src/android/accessibilityservice/AccessibilityServiceTest.java
index c65ef9a56cd8..53ba140e6aad 100644
--- a/core/tests/coretests/src/android/accessibilityservice/AccessibilityServiceTest.java
+++ b/core/tests/coretests/src/android/accessibilityservice/AccessibilityServiceTest.java
@@ -16,18 +16,40 @@
 
 package android.accessibilityservice;
 
+import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY;
+import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC;
+import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.WindowManager.LayoutParams.TYPE_ACCESSIBILITY_OVERLAY;
+import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
+
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.verify;
 
+import android.content.Context;
 import android.content.Intent;
+import android.graphics.PixelFormat;
+import android.hardware.display.DisplayManager;
+import android.hardware.display.VirtualDisplay;
+import android.media.ImageReader;
+import android.os.Binder;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.RemoteException;
+import android.util.SparseArray;
+import android.view.Display;
+import android.view.View;
+import android.view.WindowManager;
+import android.view.WindowManagerGlobal;
 import android.view.accessibility.AccessibilityEvent;
+import android.window.WindowTokenClient;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.core.app.ApplicationProvider;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -42,6 +64,8 @@ import org.mockito.MockitoAnnotations;
 public class AccessibilityServiceTest {
     private static final String TAG = "AccessibilityServiceTest";
     private static final int CONNECTION_ID = 1;
+    private final WindowManager.LayoutParams mParams = new WindowManager.LayoutParams(
+            TYPE_ACCESSIBILITY_OVERLAY);
 
     private static class AccessibilityServiceTestClass extends AccessibilityService {
         private IAccessibilityServiceClient mCallback;
@@ -49,7 +73,11 @@ public class AccessibilityServiceTest {
 
         AccessibilityServiceTestClass() {
             super();
-            attachBaseContext(InstrumentationRegistry.getContext());
+            Context context = ApplicationProvider.getApplicationContext();
+            final Display display = context.getSystemService(DisplayManager.class)
+                    .getDisplay(DEFAULT_DISPLAY);
+
+            attachBaseContext(context.createTokenContext(new WindowTokenClient(), display));
             mLooper = InstrumentationRegistry.getContext().getMainLooper();
         }
 
@@ -78,14 +106,33 @@ public class AccessibilityServiceTest {
     private @Mock IBinder mMockIBinder;
     private IAccessibilityServiceClient mServiceInterface;
     private AccessibilityServiceTestClass mService;
+    private final SparseArray<IBinder> mWindowTokens = new SparseArray<>();
 
     @Before
-    public void setUp() throws RemoteException {
+    public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         mService = new AccessibilityServiceTestClass();
+        mService.onCreate();
         mService.setupCallback(mMockClientForCallback);
         mServiceInterface = (IAccessibilityServiceClient) mService.onBind(new Intent());
         mServiceInterface.init(mMockConnection, CONNECTION_ID, mMockIBinder);
+        doAnswer(invocation -> {
+            Object[] args = invocation.getArguments();
+            final int displayId = (int) args[0];
+            final IBinder token = new Binder();
+            WindowManagerGlobal.getWindowManagerService().addWindowToken(token,
+                    TYPE_ACCESSIBILITY_OVERLAY, displayId, null /* options */);
+            mWindowTokens.put(displayId, token);
+            return token;
+        }).when(mMockConnection).getOverlayWindowToken(anyInt());
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        for (int i = mWindowTokens.size() - 1; i >= 0; --i) {
+            WindowManagerGlobal.getWindowManagerService().removeWindowToken(
+                    mWindowTokens.valueAt(i), mWindowTokens.keyAt(i));
+        }
     }
 
     @Test
@@ -101,4 +148,79 @@ public class AccessibilityServiceTest {
 
         verify(mMockConnection).getSystemActions();
     }
+
+    @Test
+    public void testAddViewWithA11yServiceDerivedDisplayContext() throws Exception {
+        try (VirtualDisplaySession session = new VirtualDisplaySession()) {
+            final Context context = mService.createDisplayContext(session.getDisplay());
+            InstrumentationRegistry.getInstrumentation().runOnMainSync(
+                    () -> context.getSystemService(WindowManager.class)
+                            .addView(new View(context), mParams)
+            );
+        }
+    }
+
+    @Test
+    public void testAddViewWithA11yServiceDerivedWindowContext() throws Exception {
+        try (VirtualDisplaySession session = new VirtualDisplaySession()) {
+            final Context context = mService.createDisplayContext(session.getDisplay())
+                    .createWindowContext(TYPE_ACCESSIBILITY_OVERLAY, null /* options */);
+            InstrumentationRegistry.getInstrumentation().runOnMainSync(
+                    () -> context.getSystemService(WindowManager.class)
+                            .addView(new View(context), mParams)
+            );
+        }
+    }
+
+    @Test
+    public void testAddViewWithA11yServiceDerivedWindowContextWithDisplay() throws Exception {
+        try (VirtualDisplaySession session = new VirtualDisplaySession()) {
+            final Context context = mService.createWindowContext(session.getDisplay(),
+                    TYPE_ACCESSIBILITY_OVERLAY, null /* options */);
+            InstrumentationRegistry.getInstrumentation().runOnMainSync(
+                    () -> context.getSystemService(WindowManager.class)
+                            .addView(new View(context), mParams)
+            );
+        }
+    }
+
+    @Test(expected = WindowManager.BadTokenException.class)
+    public void testAddViewWithA11yServiceDerivedWindowContextWithDifferentType()
+            throws Exception {
+        try (VirtualDisplaySession session = new VirtualDisplaySession()) {
+            final Context context = mService.createWindowContext(session.getDisplay(),
+                    TYPE_APPLICATION_OVERLAY, null /* options */);
+            InstrumentationRegistry.getInstrumentation().runOnMainSync(
+                    () -> context.getSystemService(WindowManager.class)
+                            .addView(new View(context), mParams)
+            );
+        }
+    }
+
+
+    private static class VirtualDisplaySession implements AutoCloseable {
+        private final VirtualDisplay mVirtualDisplay;
+
+        VirtualDisplaySession() {
+            final DisplayManager displayManager = ApplicationProvider.getApplicationContext()
+                    .getSystemService(DisplayManager.class);
+            final int width = 800;
+            final int height = 480;
+            final int density = 160;
+            ImageReader reader = ImageReader.newInstance(width, height, PixelFormat.RGBA_8888,
+                    2 /* maxImages */);
+            mVirtualDisplay = displayManager.createVirtualDisplay(
+                    TAG, width, height, density, reader.getSurface(),
+                    VIRTUAL_DISPLAY_FLAG_PUBLIC | VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY);
+        }
+
+        private Display getDisplay() {
+            return mVirtualDisplay.getDisplay();
+        }
+
+        @Override
+        public void close() throws Exception {
+            mVirtualDisplay.release();
+        }
+    }
 }
diff --git a/core/tests/coretests/src/android/app/NotificationTest.java b/core/tests/coretests/src/android/app/NotificationTest.java
index 34c1763b3286..37cf514e92ea 100644
--- a/core/tests/coretests/src/android/app/NotificationTest.java
+++ b/core/tests/coretests/src/android/app/NotificationTest.java
@@ -43,6 +43,7 @@ import android.graphics.BitmapFactory;
 import android.graphics.Color;
 import android.graphics.drawable.Icon;
 import android.os.Build;
+import android.os.Bundle;
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.text.Spannable;
@@ -545,6 +546,29 @@ public class NotificationTest {
         validateColorizedPaletteForColor(Color.BLACK);
     }
 
+    @Test
+    public void testIsMediaNotification_nullSession_returnsFalse() {
+        // Null media session
+        Notification.MediaStyle mediaStyle = new Notification.MediaStyle();
+        Notification notification = new Notification.Builder(mContext, "test id")
+                .setStyle(mediaStyle)
+                .build();
+        assertFalse(notification.isMediaNotification());
+    }
+
+    @Test
+    public void testIsMediaNotification_invalidSession_returnsFalse() {
+        // Extra was set manually to an invalid type
+        Bundle extras = new Bundle();
+        extras.putParcelable(Notification.EXTRA_MEDIA_SESSION, new Intent());
+        Notification.MediaStyle mediaStyle = new Notification.MediaStyle();
+        Notification notification = new Notification.Builder(mContext, "test id")
+                .setStyle(mediaStyle)
+                .addExtras(extras)
+                .build();
+        assertFalse(notification.isMediaNotification());
+    }
+
     public void validateColorizedPaletteForColor(int rawColor) {
         Notification.Colors cDay = new Notification.Colors();
         Notification.Colors cNight = new Notification.Colors();
diff --git a/core/tests/coretests/src/android/app/activity/ActivityThreadTest.java b/core/tests/coretests/src/android/app/activity/ActivityThreadTest.java
index fb820cb2f5e5..4a7c50d8b934 100644
--- a/core/tests/coretests/src/android/app/activity/ActivityThreadTest.java
+++ b/core/tests/coretests/src/android/app/activity/ActivityThreadTest.java
@@ -286,7 +286,6 @@ public class ActivityThreadTest {
     }
 
     @Test
-    @FlakyTest(bugId = 194242735)
     public void testHandleActivityConfigurationChanged_EnsureUpdatesProcessedInOrder()
             throws Exception {
         final TestActivity activity = mActivityTestRule.launchActivity(new Intent());
@@ -305,18 +304,22 @@ public class ActivityThreadTest {
         final int numOfConfig = activity.mNumOfConfigChanges;
 
         final Configuration processConfigLandscape = new Configuration();
+        processConfigLandscape.orientation = ORIENTATION_LANDSCAPE;
         processConfigLandscape.windowConfiguration.setBounds(new Rect(0, 0, 100, 60));
         processConfigLandscape.seq = BASE_SEQ + 1;
 
         final Configuration activityConfigLandscape = new Configuration();
+        activityConfigLandscape.orientation = ORIENTATION_LANDSCAPE;
         activityConfigLandscape.windowConfiguration.setBounds(new Rect(0, 0, 100, 50));
         activityConfigLandscape.seq = BASE_SEQ + 2;
 
         final Configuration processConfigPortrait = new Configuration();
+        processConfigPortrait.orientation = ORIENTATION_PORTRAIT;
         processConfigPortrait.windowConfiguration.setBounds(new Rect(0, 0, 60, 100));
         processConfigPortrait.seq = BASE_SEQ + 3;
 
         final Configuration activityConfigPortrait = new Configuration();
+        activityConfigPortrait.orientation = ORIENTATION_PORTRAIT;
         activityConfigPortrait.windowConfiguration.setBounds(new Rect(0, 0, 50, 100));
         activityConfigPortrait.seq = BASE_SEQ + 4;
 
@@ -344,7 +347,8 @@ public class ActivityThreadTest {
         assertEquals(activityConfigPortrait.windowConfiguration.getBounds(), bounds);
 
         // Ensure that Activity#onConfigurationChanged() not be called because the changes in
-        // WindowConfiguration shouldn't be reported.
+        // WindowConfiguration shouldn't be reported, and we only apply the latest Configuration
+        // update in transaction.
         assertEquals(numOfConfig, activity.mNumOfConfigChanges);
     }
 
diff --git a/core/tests/coretests/src/android/app/appsearch/external/app/GenericDocumentTest.java b/core/tests/coretests/src/android/app/appsearch/external/app/GenericDocumentTest.java
index 3d820acf2d22..6884f13d4cc9 100644
--- a/core/tests/coretests/src/android/app/appsearch/external/app/GenericDocumentTest.java
+++ b/core/tests/coretests/src/android/app/appsearch/external/app/GenericDocumentTest.java
@@ -62,18 +62,4 @@ public class GenericDocumentTest {
         assertThat(outDoc.getPropertyDocument("propDocument").getPropertyBytesArray("propBytes"))
                 .isEqualTo(new byte[][] {{3, 4}});
     }
-
-    @Test
-    public void testPutLargeDocument_exceedLimit() throws Exception {
-        // Create a String property that has a very large property.
-        char[] chars = new char[10_000_000];
-        String property = new StringBuilder().append(chars).append("the end.").toString();
-
-        GenericDocument doc =
-                new GenericDocument.Builder<>("namespace", "id1", "schema1")
-                        .setPropertyString("propString", property)
-                        .build();
-
-        assertThat(doc.getPropertyString("propString")).isEqualTo(property);
-    }
 }
diff --git a/core/tests/coretests/src/android/app/servertransaction/ObjectPoolTests.java b/core/tests/coretests/src/android/app/servertransaction/ObjectPoolTests.java
index 9915e3852b8d..3e261a7113ac 100644
--- a/core/tests/coretests/src/android/app/servertransaction/ObjectPoolTests.java
+++ b/core/tests/coretests/src/android/app/servertransaction/ObjectPoolTests.java
@@ -156,7 +156,8 @@ public class ObjectPoolTests {
                 .setProcState(procState).setState(bundle).setPersistentState(persistableBundle)
                 .setPendingResults(resultInfoList()).setPendingNewIntents(referrerIntentList())
                 .setIsForward(true).setAssistToken(assistToken)
-                .setShareableActivityToken(shareableActivityToken).build();
+                .setShareableActivityToken(shareableActivityToken)
+                .build();
 
         LaunchActivityItem emptyItem = new LaunchActivityItemBuilder().build();
         LaunchActivityItem item = itemSupplier.get();
diff --git a/core/tests/coretests/src/android/app/servertransaction/TransactionParcelTests.java b/core/tests/coretests/src/android/app/servertransaction/TransactionParcelTests.java
index df0c64c810c6..9f48c06da385 100644
--- a/core/tests/coretests/src/android/app/servertransaction/TransactionParcelTests.java
+++ b/core/tests/coretests/src/android/app/servertransaction/TransactionParcelTests.java
@@ -532,7 +532,7 @@ public class TransactionParcelTests {
         }
 
         @Override
-        public void scheduleCrash(String s, int i) throws RemoteException {
+        public void scheduleCrash(String s, int i, Bundle extras) throws RemoteException {
         }
 
         @Override
diff --git a/core/tests/coretests/src/android/content/ContextTest.java b/core/tests/coretests/src/android/content/ContextTest.java
index d1776fb7e5c1..3d7d807ca53d 100644
--- a/core/tests/coretests/src/android/content/ContextTest.java
+++ b/core/tests/coretests/src/android/content/ContextTest.java
@@ -32,7 +32,6 @@ import android.content.res.Configuration;
 import android.graphics.PixelFormat;
 import android.hardware.display.DisplayManager;
 import android.hardware.display.VirtualDisplay;
-import android.inputmethodservice.InputMethodService;
 import android.media.ImageReader;
 import android.os.UserHandle;
 import android.view.Display;
@@ -140,13 +139,6 @@ public class ContextTest {
     }
 
     @Test
-    public void testIsUiContext_InputMethodService_returnsTrue() {
-        final InputMethodService ims = new InputMethodService();
-
-        assertTrue(ims.isUiContext());
-    }
-
-    @Test
     public void testGetDisplayFromDisplayContextDerivedContextOnPrimaryDisplay() {
         verifyGetDisplayFromDisplayContextDerivedContext(false /* onSecondaryDisplay */);
     }
diff --git a/core/tests/coretests/src/android/content/pm/ConstrainDisplayApisConfigTest.java b/core/tests/coretests/src/android/content/pm/ConstrainDisplayApisConfigTest.java
index 0456029f45a0..98485c024a59 100644
--- a/core/tests/coretests/src/android/content/pm/ConstrainDisplayApisConfigTest.java
+++ b/core/tests/coretests/src/android/content/pm/ConstrainDisplayApisConfigTest.java
@@ -18,8 +18,7 @@ package android.content.pm;
 
 import static android.provider.DeviceConfig.NAMESPACE_CONSTRAIN_DISPLAY_APIS;
 
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertEquals;
 
 import android.annotation.Nullable;
 import android.platform.test.annotations.Presubmit;
@@ -146,24 +145,17 @@ public final class ConstrainDisplayApisConfigTest {
 
     private static void testNeverConstrainDisplayApis(String packageName, long version,
             boolean expected) {
-        boolean result = ConstrainDisplayApisConfig.neverConstrainDisplayApis(
-                buildApplicationInfo(packageName, version));
-        if (expected) {
-            assertTrue(result);
-        } else {
-            assertFalse(result);
-        }
+        ConstrainDisplayApisConfig config = new ConstrainDisplayApisConfig();
+        assertEquals(expected,
+                config.getNeverConstrainDisplayApis(buildApplicationInfo(packageName, version)));
     }
 
     private static void testAlwaysConstrainDisplayApis(String packageName, long version,
             boolean expected) {
-        boolean result = ConstrainDisplayApisConfig.alwaysConstrainDisplayApis(
-                buildApplicationInfo(packageName, version));
-        if (expected) {
-            assertTrue(result);
-        } else {
-            assertFalse(result);
-        }
+        ConstrainDisplayApisConfig config = new ConstrainDisplayApisConfig();
+
+        assertEquals(expected,
+                config.getAlwaysConstrainDisplayApis(buildApplicationInfo(packageName, version)));
     }
 
     private static ApplicationInfo buildApplicationInfo(String packageName, long version) {
diff --git a/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java b/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
index 8fd1af801094..4f1da1b29616 100644
--- a/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
+++ b/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
@@ -111,7 +111,7 @@ public class InsetsAnimationControlImplTest {
                 new Rect(0, 0, 500, 500), mInsetsState, mMockListener, systemBars(),
                 mMockController, 10 /* durationMs */, new LinearInterpolator(),
                 0 /* animationType */, 0 /* layoutInsetsDuringAnimation */, null /* translator */);
-        mController.mReadyDispatched = true;
+        mController.setReadyDispatched(true);
     }
 
     @Test
@@ -197,7 +197,7 @@ public class InsetsAnimationControlImplTest {
 
     @Test
     public void testCancelled_beforeReadyDispatched() {
-        mController.mReadyDispatched = false;
+        mController.setReadyDispatched(false);
         mController.cancel();
         assertFalse(mController.isReady());
         assertFalse(mController.isFinished());
diff --git a/core/tests/coretests/src/android/view/InsetsControllerTest.java b/core/tests/coretests/src/android/view/InsetsControllerTest.java
index 6301f32169f7..227a86576113 100644
--- a/core/tests/coretests/src/android/view/InsetsControllerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsControllerTest.java
@@ -19,13 +19,17 @@ package android.view;
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
 import static android.view.InsetsController.ANIMATION_TYPE_HIDE;
 import static android.view.InsetsController.ANIMATION_TYPE_NONE;
+import static android.view.InsetsController.ANIMATION_TYPE_RESIZE;
 import static android.view.InsetsController.ANIMATION_TYPE_SHOW;
+import static android.view.InsetsController.AnimationType;
 import static android.view.InsetsSourceConsumer.ShowResult.IME_SHOW_DELAYED;
 import static android.view.InsetsSourceConsumer.ShowResult.SHOW_IMMEDIATELY;
+import static android.view.InsetsState.FIRST_TYPE;
 import static android.view.InsetsState.ITYPE_CAPTION_BAR;
 import static android.view.InsetsState.ITYPE_IME;
 import static android.view.InsetsState.ITYPE_NAVIGATION_BAR;
 import static android.view.InsetsState.ITYPE_STATUS_BAR;
+import static android.view.InsetsState.LAST_TYPE;
 import static android.view.WindowInsets.Type.ime;
 import static android.view.WindowInsets.Type.navigationBars;
 import static android.view.WindowInsets.Type.statusBars;
@@ -662,6 +666,97 @@ public class InsetsControllerTest {
     }
 
     @Test
+    public void testResizeAnimation_insetsTypes() {
+        for (@InternalInsetsType int type = FIRST_TYPE; type <= LAST_TYPE; type++) {
+            final @AnimationType int expectedAnimationType =
+                    (InsetsState.toPublicType(type) & Type.systemBars()) != 0
+                            ? ANIMATION_TYPE_RESIZE
+                            : ANIMATION_TYPE_NONE;
+            doTestResizeAnimation_insetsTypes(type, expectedAnimationType);
+        }
+    }
+
+    private void doTestResizeAnimation_insetsTypes(@InternalInsetsType int type,
+            @AnimationType int expectedAnimationType) {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            final InsetsState state1 = new InsetsState();
+            state1.getSource(type).setVisible(true);
+            state1.getSource(type).setFrame(0, 0, 500, 50);
+            final InsetsState state2 = new InsetsState(state1, true /* copySources */);
+            state2.getSource(type).setFrame(0, 0, 500, 60);
+            final String message = "Animation type of " + InsetsState.typeToString(type) + ":";
+
+            // New insets source won't cause the resize animation.
+            mController.onStateChanged(state1);
+            assertEquals(message, ANIMATION_TYPE_NONE, mController.getAnimationType(type));
+
+            // Changing frame might cause the resize animation. This depends on the insets type.
+            mController.onStateChanged(state2);
+            assertEquals(message, expectedAnimationType, mController.getAnimationType(type));
+
+            // Cancel the existing animations for the next iteration.
+            mController.cancelExistingAnimations();
+            assertEquals(message, ANIMATION_TYPE_NONE, mController.getAnimationType(type));
+        });
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+    }
+
+    @Test
+    public void testResizeAnimation_displayFrame() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            final @InternalInsetsType int type = ITYPE_STATUS_BAR;
+            final InsetsState state1 = new InsetsState();
+            state1.setDisplayFrame(new Rect(0, 0, 500, 1000));
+            state1.getSource(type).setFrame(0, 0, 500, 50);
+            final InsetsState state2 = new InsetsState(state1, true /* copySources */);
+            state2.setDisplayFrame(new Rect(0, 0, 500, 1010));
+            state2.getSource(type).setFrame(0, 0, 500, 60);
+            final String message = "There must not be resize animation.";
+
+            // New insets source won't cause the resize animation.
+            mController.onStateChanged(state1);
+            assertEquals(message, ANIMATION_TYPE_NONE, mController.getAnimationType(type));
+
+            // Changing frame won't cause the resize animation if the display frame is also changed.
+            mController.onStateChanged(state2);
+            assertEquals(message, ANIMATION_TYPE_NONE, mController.getAnimationType(type));
+        });
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+    }
+
+    @Test
+    public void testResizeAnimation_visibility() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            final @InternalInsetsType int type = ITYPE_STATUS_BAR;
+            final InsetsState state1 = new InsetsState();
+            state1.getSource(type).setVisible(true);
+            state1.getSource(type).setFrame(0, 0, 500, 50);
+            final InsetsState state2 = new InsetsState(state1, true /* copySources */);
+            state2.getSource(type).setVisible(false);
+            state2.getSource(type).setFrame(0, 0, 500, 60);
+            final InsetsState state3 = new InsetsState(state2, true /* copySources */);
+            state3.getSource(type).setVisible(true);
+            state3.getSource(type).setFrame(0, 0, 500, 70);
+            final String message = "There must not be resize animation.";
+
+            // New insets source won't cause the resize animation.
+            mController.onStateChanged(state1);
+            assertEquals(message, ANIMATION_TYPE_NONE, mController.getAnimationType(type));
+
+            // Changing source visibility (visible --> invisible) won't cause the resize animation.
+            // The previous source and the current one must be both visible.
+            mController.onStateChanged(state2);
+            assertEquals(message, ANIMATION_TYPE_NONE, mController.getAnimationType(type));
+
+            // Changing source visibility (invisible --> visible) won't cause the resize animation.
+            // The previous source and the current one must be both visible.
+            mController.onStateChanged(state3);
+            assertEquals(message, ANIMATION_TYPE_NONE, mController.getAnimationType(type));
+        });
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+    }
+
+    @Test
     public void testCaptionInsetsStateAssemble() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
             mController.onFrameChanged(new Rect(0, 0, 100, 300));
@@ -698,15 +793,15 @@ public class InsetsControllerTest {
     @Test
     public void testRequestedState() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
-            final InsetsState state = mTestHost.getRequestedState();
+            final InsetsVisibilities request = mTestHost.getRequestedVisibilities();
 
             mController.hide(statusBars() | navigationBars());
-            assertFalse(state.getSourceOrDefaultVisibility(ITYPE_STATUS_BAR));
-            assertFalse(state.getSourceOrDefaultVisibility(ITYPE_NAVIGATION_BAR));
+            assertFalse(request.getVisibility(ITYPE_STATUS_BAR));
+            assertFalse(request.getVisibility(ITYPE_NAVIGATION_BAR));
 
             mController.show(statusBars() | navigationBars());
-            assertTrue(state.getSourceOrDefaultVisibility(ITYPE_STATUS_BAR));
-            assertTrue(state.getSourceOrDefaultVisibility(ITYPE_NAVIGATION_BAR));
+            assertTrue(request.getVisibility(ITYPE_STATUS_BAR));
+            assertTrue(request.getVisibility(ITYPE_NAVIGATION_BAR));
         });
     }
 
@@ -837,20 +932,20 @@ public class InsetsControllerTest {
 
     public static class TestHost extends ViewRootInsetsControllerHost {
 
-        private final InsetsState mRequestedState = new InsetsState();
+        private final InsetsVisibilities mRequestedVisibilities = new InsetsVisibilities();
 
         TestHost(ViewRootImpl viewRoot) {
             super(viewRoot);
         }
 
         @Override
-        public void onInsetsModified(InsetsState insetsState) {
-            mRequestedState.set(insetsState, true);
-            super.onInsetsModified(insetsState);
+        public void updateRequestedVisibilities(InsetsVisibilities visibilities) {
+            mRequestedVisibilities.set(visibilities);
+            super.updateRequestedVisibilities(visibilities);
         }
 
-        public InsetsState getRequestedState() {
-            return mRequestedState;
+        public InsetsVisibilities getRequestedVisibilities() {
+            return mRequestedVisibilities;
         }
     }
 }
diff --git a/core/tests/coretests/src/android/view/InsetsStateTest.java b/core/tests/coretests/src/android/view/InsetsStateTest.java
index 3bd29398325f..bf8bb76891d7 100644
--- a/core/tests/coretests/src/android/view/InsetsStateTest.java
+++ b/core/tests/coretests/src/android/view/InsetsStateTest.java
@@ -216,9 +216,9 @@ public class InsetsStateTest {
         mState.getSource(ITYPE_CAPTION_BAR).setFrame(new Rect(0, 0, 100, 300));
         mState.getSource(ITYPE_CAPTION_BAR).setVisible(true);
 
-        Rect visibleInsets = mState.calculateVisibleInsets(
+        Insets visibleInsets = mState.calculateVisibleInsets(
                 new Rect(0, 0, 100, 400), SOFT_INPUT_ADJUST_NOTHING);
-        assertEquals(new Rect(0, 300, 0, 0), visibleInsets);
+        assertEquals(Insets.of(0, 300, 0, 0), visibleInsets);
     }
 
     @Test
@@ -226,9 +226,9 @@ public class InsetsStateTest {
         mState.getSource(ITYPE_CAPTION_BAR).setFrame(new Rect(0, 0, 100, 300));
         mState.getSource(ITYPE_CAPTION_BAR).setVisible(true);
 
-        Rect visibleInsets = mState.calculateVisibleInsets(
+        Insets visibleInsets = mState.calculateVisibleInsets(
                 new Rect(0, 0, 150, 400), SOFT_INPUT_ADJUST_NOTHING);
-        assertEquals(new Rect(0, 300, 0, 0), visibleInsets);
+        assertEquals(Insets.of(0, 300, 0, 0), visibleInsets);
     }
 
     @Test
@@ -413,9 +413,9 @@ public class InsetsStateTest {
         // Make sure bottom gestures are ignored
         mState.getSource(ITYPE_BOTTOM_GESTURES).setFrame(new Rect(0, 100, 100, 300));
         mState.getSource(ITYPE_BOTTOM_GESTURES).setVisible(true);
-        Rect visibleInsets = mState.calculateVisibleInsets(
+        Insets visibleInsets = mState.calculateVisibleInsets(
                 new Rect(0, 0, 100, 300), SOFT_INPUT_ADJUST_PAN);
-        assertEquals(new Rect(0, 100, 0, 100), visibleInsets);
+        assertEquals(Insets.of(0, 100, 0, 100), visibleInsets);
     }
 
     @Test
@@ -428,9 +428,9 @@ public class InsetsStateTest {
         // Make sure bottom gestures are ignored
         mState.getSource(ITYPE_BOTTOM_GESTURES).setFrame(new Rect(0, 100, 100, 300));
         mState.getSource(ITYPE_BOTTOM_GESTURES).setVisible(true);
-        Rect visibleInsets = mState.calculateVisibleInsets(
+        Insets visibleInsets = mState.calculateVisibleInsets(
                 new Rect(0, 0, 100, 300), SOFT_INPUT_ADJUST_NOTHING);
-        assertEquals(new Rect(0, 100, 0, 0), visibleInsets);
+        assertEquals(Insets.of(0, 100, 0, 0), visibleInsets);
     }
 
     @Test
diff --git a/core/tests/coretests/src/android/view/InsetsVisibilitiesTest.java b/core/tests/coretests/src/android/view/InsetsVisibilitiesTest.java
new file mode 100644
index 000000000000..5664e0b0aa0f
--- /dev/null
+++ b/core/tests/coretests/src/android/view/InsetsVisibilitiesTest.java
@@ -0,0 +1,121 @@
+/*
+ * Copyright (C) 2021 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 android.view;
+
+import static android.view.InsetsState.FIRST_TYPE;
+import static android.view.InsetsState.LAST_TYPE;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
+
+import android.platform.test.annotations.Presubmit;
+import android.view.InsetsState.InternalInsetsType;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for {@link InsetsVisibilities}.
+ *
+ * <p>Build/Install/Run:
+ *  atest FrameworksCoreTests:InsetsVisibilities
+ *
+ * <p>This test class is a part of Window Manager Service tests and specified in
+ * {@link com.android.server.wm.test.filters.FrameworksTestsFilter}.
+ */
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class InsetsVisibilitiesTest {
+
+    @Test
+    public void testEquals() {
+        final InsetsVisibilities v1 = new InsetsVisibilities();
+        final InsetsVisibilities v2 = new InsetsVisibilities();
+        final InsetsVisibilities v3 = new InsetsVisibilities();
+        assertEquals(v1, v2);
+        assertEquals(v1, v3);
+
+        for (@InternalInsetsType int type = FIRST_TYPE; type <= LAST_TYPE; type++) {
+            v1.setVisibility(type, false);
+            v2.setVisibility(type, false);
+        }
+        assertEquals(v1, v2);
+        assertNotEquals(v1, v3);
+
+        for (@InternalInsetsType int type = FIRST_TYPE; type <= LAST_TYPE; type++) {
+            v1.setVisibility(type, true);
+            v2.setVisibility(type, true);
+        }
+        assertEquals(v1, v2);
+        assertNotEquals(v1, v3);
+    }
+
+    @Test
+    public void testSet() {
+        for (@InternalInsetsType int type = FIRST_TYPE; type <= LAST_TYPE; type++) {
+            final InsetsVisibilities v1 = new InsetsVisibilities();
+            final InsetsVisibilities v2 = new InsetsVisibilities();
+
+            v1.setVisibility(type, true);
+            assertNotEquals(v1, v2);
+
+            v2.set(v1);
+            assertEquals(v1, v2);
+
+            v2.setVisibility(type, false);
+            assertNotEquals(v1, v2);
+
+            v1.set(v2);
+            assertEquals(v1, v2);
+        }
+    }
+
+    @Test
+    public void testCopyConstructor() {
+        for (@InternalInsetsType int type = FIRST_TYPE; type <= LAST_TYPE; type++) {
+            final InsetsVisibilities v1 = new InsetsVisibilities();
+            v1.setVisibility(type, true);
+            final InsetsVisibilities v2 = new InsetsVisibilities(v1);
+            assertEquals(v1, v2);
+
+            v2.setVisibility(type, false);
+            assertNotEquals(v1, v2);
+        }
+    }
+
+    @Test
+    public void testGetterAndSetter() {
+        final InsetsVisibilities v1 = new InsetsVisibilities();
+        final InsetsVisibilities v2 = new InsetsVisibilities();
+
+        for (@InternalInsetsType int type = FIRST_TYPE; type <= LAST_TYPE; type++) {
+            assertEquals(InsetsState.getDefaultVisibility(type), v1.getVisibility(type));
+        }
+
+        for (@InternalInsetsType int type = FIRST_TYPE; type <= LAST_TYPE; type++) {
+            v1.setVisibility(type, true);
+            assertTrue(v1.getVisibility(type));
+
+            v2.setVisibility(type, false);
+            assertFalse(v2.getVisibility(type));
+        }
+    }
+}
diff --git a/core/tests/coretests/src/android/view/MotionEventTest.java b/core/tests/coretests/src/android/view/MotionEventTest.java
index b3450de80092..b6a182c8faf4 100644
--- a/core/tests/coretests/src/android/view/MotionEventTest.java
+++ b/core/tests/coretests/src/android/view/MotionEventTest.java
@@ -16,6 +16,7 @@
 
 package android.view;
 
+import static android.view.InputDevice.SOURCE_CLASS_POINTER;
 import static android.view.MotionEvent.ACTION_DOWN;
 import static android.view.MotionEvent.ACTION_POINTER_DOWN;
 import static android.view.MotionEvent.TOOL_TYPE_FINGER;
@@ -189,4 +190,27 @@ public class MotionEventTest {
         assertEquals(950, (int) rot270.getX());
         assertEquals(30, (int) rot270.getY());
     }
+
+    @Test
+    public void testUsesPointerSourceByDefault() {
+        final MotionEvent event = MotionEvent.obtain(0 /* downTime */, 0 /* eventTime */,
+                ACTION_DOWN, 0 /* x */, 0 /* y */, 0 /* metaState */);
+        assertTrue(event.isFromSource(SOURCE_CLASS_POINTER));
+    }
+
+    @Test
+    public void testLocationOffsetOnlyAppliedToNonPointerSources() {
+        final MotionEvent event = MotionEvent.obtain(0 /* downTime */, 0 /* eventTime */,
+                ACTION_DOWN, 10 /* x */, 20 /* y */, 0 /* metaState */);
+        event.offsetLocation(40, 50);
+
+        // The offset should be applied since a pointer source is used by default.
+        assertEquals(50, (int) event.getX());
+        assertEquals(70, (int) event.getY());
+
+        // The offset should not be applied if the source is changed to a non-pointer source.
+        event.setSource(InputDevice.SOURCE_JOYSTICK);
+        assertEquals(10, (int) event.getX());
+        assertEquals(20, (int) event.getY());
+    }
 }
diff --git a/core/tests/coretests/src/android/view/WindowInfoTest.java b/core/tests/coretests/src/android/view/WindowInfoTest.java
index 05e8bd8b6cab..0a99b08f58ff 100644
--- a/core/tests/coretests/src/android/view/WindowInfoTest.java
+++ b/core/tests/coretests/src/android/view/WindowInfoTest.java
@@ -25,6 +25,7 @@ import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 
+import android.app.ActivityTaskManager;
 import android.os.IBinder;
 import android.os.Parcel;
 import android.platform.test.annotations.Presubmit;
@@ -83,6 +84,7 @@ public class WindowInfoTest {
         assertEquals(0, w.layer);
         assertEquals(AccessibilityNodeInfo.UNDEFINED_NODE_ID, w.accessibilityIdOfAnchor);
         assertEquals(Display.INVALID_DISPLAY, w.displayId);
+        assertEquals(ActivityTaskManager.INVALID_TASK_ID, w.taskId);
         assertNull(w.title);
         assertNull(w.token);
         assertNull(w.childTokens);
@@ -123,6 +125,7 @@ public class WindowInfoTest {
         windowInfo.displayId = 2;
         windowInfo.layer = 3;
         windowInfo.accessibilityIdOfAnchor = 4L;
+        windowInfo.taskId = 5;
         windowInfo.title = "title";
         windowInfo.token = mock(IBinder.class);
         windowInfo.childTokens = new ArrayList<>();
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
index 8643a37bba8d..b71d814c508d 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityServiceConnectionImpl.java
@@ -166,4 +166,7 @@ public class AccessibilityServiceConnectionImpl extends IAccessibilityServiceCon
 
     public void logTrace(long timestamp, String where, String callingParams, int processId,
             long threadId, int callingUid, Bundle callingStack) {}
+
+    public void logTrace(long timestamp, String where, long loggingTypes, String callingParams,
+            int processId, long threadId, int callingUid, Bundle serializedCallingStackInBundle) {}
 }
diff --git a/core/tests/coretests/src/android/view/contentcapture/ContentCaptureContextTest.java b/core/tests/coretests/src/android/view/contentcapture/ContentCaptureContextTest.java
index ddb6729b55e1..4b19391da78e 100644
--- a/core/tests/coretests/src/android/view/contentcapture/ContentCaptureContextTest.java
+++ b/core/tests/coretests/src/android/view/contentcapture/ContentCaptureContextTest.java
@@ -39,9 +39,10 @@ public class ContentCaptureContextTest {
     public void testConstructorAdditionalFlags() {
         final ComponentName componentName = new ComponentName("component", "name");
         final IBinder token = new Binder();
+        final IBinder windowToken = new Binder();
         final ContentCaptureContext ctx = new ContentCaptureContext(/* clientContext= */ null,
                 new ActivityId(/* taskId= */ 666, token), componentName, /* displayId= */
-                42, /* flags= */ 1);
+                42, windowToken, /* flags= */ 1);
         final ContentCaptureContext newCtx = new ContentCaptureContext(ctx, /* extraFlags= */ 2);
         assertThat(newCtx.getFlags()).isEqualTo(3);
         assertThat(newCtx.getActivityComponent()).isEqualTo(componentName);
@@ -50,6 +51,7 @@ public class ContentCaptureContextTest {
         assertThat(activityId.getTaskId()).isEqualTo(666);
         assertThat(activityId.getToken()).isEqualTo(token);
         assertThat(newCtx.getDisplayId()).isEqualTo(42);
+        assertThat(newCtx.getWindowToken()).isEqualTo(windowToken);
         assertThat(newCtx.getExtras()).isNull();
         assertThat(newCtx.getLocusId()).isNull();
         assertThat(newCtx.getParentSessionId()).isNull();
diff --git a/core/tests/coretests/src/android/window/CustomDrawable.java b/core/tests/coretests/src/android/window/CustomDrawable.java
new file mode 100644
index 000000000000..c25f87782420
--- /dev/null
+++ b/core/tests/coretests/src/android/window/CustomDrawable.java
@@ -0,0 +1,25 @@
+/*
+ * Copyright (C) 2021 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 android.window;
+
+import android.graphics.drawable.InsetDrawable;
+
+public class CustomDrawable extends InsetDrawable {
+    public CustomDrawable() {
+        super(null, 0);
+    }
+}
diff --git a/core/tests/coretests/src/android/window/WindowContextControllerTest.java b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
index d7b36a8d7067..52cb9f318dd0 100644
--- a/core/tests/coretests/src/android/window/WindowContextControllerTest.java
+++ b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
@@ -22,6 +22,7 @@ import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
@@ -61,7 +62,7 @@ public class WindowContextControllerTest {
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         mController = new WindowContextController(mMockToken);
-        doNothing().when(mMockToken).onConfigurationChanged(any(), anyInt());
+        doNothing().when(mMockToken).onConfigurationChanged(any(), anyInt(), anyBoolean());
         doReturn(true).when(mMockToken).attachToDisplayArea(anyInt(), anyInt(), any());
     }
 
diff --git a/core/tests/coretests/src/android/window/WindowContextTest.java b/core/tests/coretests/src/android/window/WindowContextTest.java
index 83280f18c889..656e756416d0 100644
--- a/core/tests/coretests/src/android/window/WindowContextTest.java
+++ b/core/tests/coretests/src/android/window/WindowContextTest.java
@@ -23,6 +23,7 @@ import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
 import android.app.Activity;
@@ -47,6 +48,8 @@ import androidx.test.platform.app.InstrumentationRegistry;
 import androidx.test.rule.ActivityTestRule;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.frameworks.coretests.R;
+
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -242,6 +245,12 @@ public class WindowContextTest {
         mInstrumentation.runOnMainSync(() -> wm.addView(subWindow, subWindowAttrs));
     }
 
+    @Test
+    public void testGetCustomDrawable() {
+        assertNotNull(mWindowContext.getResources().getDrawable(R.drawable.custom_drawable,
+                null /* theme */));
+    }
+
     private WindowContext createWindowContext() {
         return createWindowContext(TYPE_APPLICATION_OVERLAY);
     }
diff --git a/core/tests/coretests/src/com/android/internal/jank/FrameTrackerTest.java b/core/tests/coretests/src/com/android/internal/jank/FrameTrackerTest.java
index 96b4316ffafc..7cd8197ce1e4 100644
--- a/core/tests/coretests/src/com/android/internal/jank/FrameTrackerTest.java
+++ b/core/tests/coretests/src/com/android/internal/jank/FrameTrackerTest.java
@@ -23,6 +23,7 @@ import static android.view.SurfaceControl.JankData.JANK_SURFACEFLINGER_DEADLINE_
 import static com.android.internal.jank.FrameTracker.SurfaceControlWrapper;
 import static com.android.internal.jank.FrameTracker.ViewRootWrapper;
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE;
+import static com.android.internal.jank.InteractionJankMonitor.CUJ_WALLPAPER_TRANSITION;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -50,6 +51,7 @@ import androidx.test.rule.ActivityTestRule;
 import com.android.internal.jank.FrameTracker.ChoreographerWrapper;
 import com.android.internal.jank.FrameTracker.FrameMetricsWrapper;
 import com.android.internal.jank.FrameTracker.ThreadedRendererWrapper;
+import com.android.internal.jank.InteractionJankMonitor.Configuration;
 import com.android.internal.jank.InteractionJankMonitor.Session;
 
 import org.junit.Before;
@@ -69,7 +71,6 @@ public class FrameTrackerTest {
     public ActivityTestRule<ViewAttachTestActivity> mRule =
             new ActivityTestRule<>(ViewAttachTestActivity.class);
 
-    private FrameTracker mTracker;
     private ThreadedRendererWrapper mRenderer;
     private FrameMetricsWrapper mWrapper;
     private SurfaceControlWrapper mSurfaceControlWrapper;
@@ -85,7 +86,6 @@ public class FrameTrackerTest {
         View view = mActivity.getWindow().getDecorView();
         assertThat(view.isAttachedToWindow()).isTrue();
 
-        Handler handler = mRule.getActivity().getMainThreadHandler();
         mWrapper = Mockito.spy(new FrameMetricsWrapper());
         mRenderer = Mockito.spy(new ThreadedRendererWrapper(view.getThreadedRenderer()));
         doNothing().when(mRenderer).addObserver(any());
@@ -103,229 +103,355 @@ public class FrameTrackerTest {
                 mListenerCapture.capture());
 
         mChoreographer = mock(ChoreographerWrapper.class);
+    }
 
-        Session session = new Session(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX);
-        mTracker = Mockito.spy(
+    private FrameTracker spyFrameTracker(int cuj, String postfix, boolean surfaceOnly) {
+        Handler handler = mRule.getActivity().getMainThreadHandler();
+        Session session = new Session(cuj, postfix);
+        Configuration config = mock(Configuration.class);
+        when(config.isSurfaceOnly()).thenReturn(surfaceOnly);
+        when(config.getSurfaceControl()).thenReturn(mSurfaceControl);
+        FrameTracker frameTracker = Mockito.spy(
                 new FrameTracker(session, handler, mRenderer, mViewRootWrapper,
                         mSurfaceControlWrapper, mChoreographer, mWrapper,
-                        /*traceThresholdMissedFrames=*/ 1, /*traceThresholdFrameTimeMillis=*/ -1,
-                        null));
-        doNothing().when(mTracker).triggerPerfetto();
-        doNothing().when(mTracker).postTraceStartMarker();
+                        /* traceThresholdMissedFrames= */ 1,
+                        /* traceThresholdFrameTimeMillis= */ -1,
+                        /* FrameTrackerListener= */ null, config));
+        doNothing().when(frameTracker).triggerPerfetto();
+        doNothing().when(frameTracker).postTraceStartMarker();
+        return frameTracker;
     }
 
     @Test
     public void testOnlyFirstWindowFrameOverThreshold() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
         // Just provide current timestamp anytime mWrapper asked for VSYNC_TIMESTAMP
         when(mWrapper.getMetric(FrameMetrics.VSYNC_TIMESTAMP))
                 .then(unusedInvocation -> System.nanoTime());
 
         when(mChoreographer.getVsyncId()).thenReturn(100L);
-        mTracker.begin();
+        tracker.begin();
         verify(mRenderer, only()).addObserver(any());
 
         // send first frame with a long duration - should not be taken into account
-        sendFirstWindowFrame(100, JANK_APP_DEADLINE_MISSED, 100L);
+        sendFirstWindowFrame(tracker, 100, JANK_APP_DEADLINE_MISSED, 100L);
 
         // send another frame with a short duration - should not be considered janky
-        sendFirstWindowFrame(5, JANK_NONE, 101L);
+        sendFirstWindowFrame(tracker, 5, JANK_NONE, 101L);
 
         // end the trace session, the last janky frame is after the end() so is discarded.
         when(mChoreographer.getVsyncId()).thenReturn(102L);
-        mTracker.end(FrameTracker.REASON_END_NORMAL);
-        sendFrame(5, JANK_NONE, 102L);
-        sendFrame(500, JANK_APP_DEADLINE_MISSED, 103L);
+        tracker.end(FrameTracker.REASON_END_NORMAL);
+        sendFrame(tracker, 5, JANK_NONE, 102L);
+        sendFrame(tracker, 500, JANK_APP_DEADLINE_MISSED, 103L);
 
-        verify(mTracker).removeObservers();
-        verify(mTracker, never()).triggerPerfetto();
+        verify(tracker).removeObservers();
+        verify(tracker, never()).triggerPerfetto();
     }
 
     @Test
     public void testSfJank() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
         when(mChoreographer.getVsyncId()).thenReturn(100L);
-        mTracker.begin();
+        tracker.begin();
         verify(mRenderer, only()).addObserver(any());
 
         // send first frame - not janky
-        sendFrame(4, JANK_NONE, 100L);
+        sendFrame(tracker, 4, JANK_NONE, 100L);
 
         // send another frame - should be considered janky
-        sendFrame(40, JANK_SURFACEFLINGER_DEADLINE_MISSED, 101L);
+        sendFrame(tracker, 40, JANK_SURFACEFLINGER_DEADLINE_MISSED, 101L);
 
         // end the trace session
         when(mChoreographer.getVsyncId()).thenReturn(102L);
-        mTracker.end(FrameTracker.REASON_END_NORMAL);
-        sendFrame(4, JANK_NONE, 102L);
+        tracker.end(FrameTracker.REASON_END_NORMAL);
+        sendFrame(tracker, 4, JANK_NONE, 102L);
 
-        verify(mTracker).removeObservers();
+        verify(tracker).removeObservers();
 
         // We detected a janky frame - trigger Perfetto
-        verify(mTracker).triggerPerfetto();
+        verify(tracker).triggerPerfetto();
     }
 
     @Test
     public void testFirstFrameJankyNoTrigger() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
         when(mChoreographer.getVsyncId()).thenReturn(100L);
-        mTracker.begin();
+        tracker.begin();
         verify(mRenderer, only()).addObserver(any());
 
         // send first frame - janky
-        sendFrame(40, JANK_APP_DEADLINE_MISSED, 100L);
+        sendFrame(tracker, 40, JANK_APP_DEADLINE_MISSED, 100L);
 
         // send another frame - not jank
-        sendFrame(4, JANK_NONE, 101L);
+        sendFrame(tracker, 4, JANK_NONE, 101L);
 
         // end the trace session
         when(mChoreographer.getVsyncId()).thenReturn(102L);
-        mTracker.end(FrameTracker.REASON_END_NORMAL);
-        sendFrame(4, JANK_NONE, 102L);
+        tracker.end(FrameTracker.REASON_END_NORMAL);
+        sendFrame(tracker, 4, JANK_NONE, 102L);
 
-        verify(mTracker).removeObservers();
+        verify(tracker).removeObservers();
 
         // We detected a janky frame - trigger Perfetto
-        verify(mTracker, never()).triggerPerfetto();
+        verify(tracker, never()).triggerPerfetto();
     }
 
     @Test
     public void testOtherFrameOverThreshold() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
         when(mChoreographer.getVsyncId()).thenReturn(100L);
-        mTracker.begin();
+        tracker.begin();
         verify(mRenderer, only()).addObserver(any());
 
         // send first frame - not janky
-        sendFrame(4, JANK_NONE, 100L);
+        sendFrame(tracker, 4, JANK_NONE, 100L);
 
         // send another frame - should be considered janky
-        sendFrame(40, JANK_APP_DEADLINE_MISSED, 101L);
+        sendFrame(tracker, 40, JANK_APP_DEADLINE_MISSED, 101L);
 
         // end the trace session
         when(mChoreographer.getVsyncId()).thenReturn(102L);
-        mTracker.end(FrameTracker.REASON_END_NORMAL);
-        sendFrame(4, JANK_NONE, 102L);
+        tracker.end(FrameTracker.REASON_END_NORMAL);
+        sendFrame(tracker, 4, JANK_NONE, 102L);
 
-        verify(mTracker).removeObservers();
+        verify(tracker).removeObservers();
 
         // We detected a janky frame - trigger Perfetto
-        verify(mTracker).triggerPerfetto();
+        verify(tracker).triggerPerfetto();
     }
 
     @Test
     public void testLastFrameOverThresholdBeforeEnd() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
         when(mChoreographer.getVsyncId()).thenReturn(100L);
-        mTracker.begin();
+        tracker.begin();
         verify(mRenderer, only()).addObserver(any());
 
         // send first frame - not janky
-        sendFrame(4, JANK_NONE, 100L);
+        sendFrame(tracker, 4, JANK_NONE, 100L);
 
         // send another frame - not janky
-        sendFrame(4, JANK_NONE, 101L);
+        sendFrame(tracker, 4, JANK_NONE, 101L);
 
         // end the trace session, simulate one more valid callback came after the end call.
         when(mChoreographer.getVsyncId()).thenReturn(102L);
-        mTracker.end(FrameTracker.REASON_END_NORMAL);
-        sendFrame(50, JANK_APP_DEADLINE_MISSED, 102L);
+        tracker.end(FrameTracker.REASON_END_NORMAL);
+        sendFrame(tracker, 50, JANK_APP_DEADLINE_MISSED, 102L);
 
         // One more callback with VSYNC after the end() vsync id.
-        sendFrame(4, JANK_NONE, 103L);
+        sendFrame(tracker, 4, JANK_NONE, 103L);
 
-        verify(mTracker).removeObservers();
+        verify(tracker).removeObservers();
 
         // We detected a janky frame - trigger Perfetto
-        verify(mTracker).triggerPerfetto();
+        verify(tracker).triggerPerfetto();
     }
 
     @Test
     public void testBeginCancel() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
         when(mChoreographer.getVsyncId()).thenReturn(100L);
-        mTracker.begin();
+        tracker.begin();
         verify(mRenderer).addObserver(any());
 
         // First frame - not janky
-        sendFrame(4, JANK_NONE, 100L);
+        sendFrame(tracker, 4, JANK_NONE, 100L);
 
         // normal frame - not janky
-        sendFrame(4, JANK_NONE, 101L);
+        sendFrame(tracker, 4, JANK_NONE, 101L);
 
         // a janky frame
-        sendFrame(50, JANK_APP_DEADLINE_MISSED, 102L);
+        sendFrame(tracker, 50, JANK_APP_DEADLINE_MISSED, 102L);
 
-        mTracker.cancel(FrameTracker.REASON_CANCEL_NORMAL);
-        verify(mTracker).removeObservers();
+        tracker.cancel(FrameTracker.REASON_CANCEL_NORMAL);
+        verify(tracker).removeObservers();
         // Since the tracker has been cancelled, shouldn't trigger perfetto.
-        verify(mTracker, never()).triggerPerfetto();
+        verify(tracker, never()).triggerPerfetto();
     }
 
     @Test
     public void testCancelIfEndVsyncIdEqualsToBeginVsyncId() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
         when(mChoreographer.getVsyncId()).thenReturn(100L);
-        mTracker.begin();
+        tracker.begin();
         verify(mRenderer, only()).addObserver(any());
 
         // end the trace session
         when(mChoreographer.getVsyncId()).thenReturn(101L);
-        mTracker.end(FrameTracker.REASON_END_NORMAL);
+        tracker.end(FrameTracker.REASON_END_NORMAL);
 
         // Since the begin vsync id (101) equals to the end vsync id (101), will be treat as cancel.
-        verify(mTracker).cancel(FrameTracker.REASON_CANCEL_SAME_VSYNC);
+        verify(tracker).cancel(FrameTracker.REASON_CANCEL_SAME_VSYNC);
 
         // Observers should be removed in this case, or FrameTracker object will be leaked.
-        verify(mTracker).removeObservers();
+        verify(tracker).removeObservers();
 
         // Should never trigger Perfetto since it is a cancel.
-        verify(mTracker, never()).triggerPerfetto();
+        verify(tracker, never()).triggerPerfetto();
     }
 
     @Test
     public void testCancelIfEndVsyncIdLessThanBeginVsyncId() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
         when(mChoreographer.getVsyncId()).thenReturn(100L);
-        mTracker.begin();
+        tracker.begin();
         verify(mRenderer, only()).addObserver(any());
 
         // end the trace session at the same vsync id, end vsync id will less than the begin one.
         // Because the begin vsync id is supposed to the next frame,
-        mTracker.end(FrameTracker.REASON_END_NORMAL);
+        tracker.end(FrameTracker.REASON_END_NORMAL);
 
         // The begin vsync id (101) is larger than the end one (100), will be treat as cancel.
-        verify(mTracker).cancel(FrameTracker.REASON_CANCEL_SAME_VSYNC);
+        verify(tracker).cancel(FrameTracker.REASON_CANCEL_SAME_VSYNC);
 
         // Observers should be removed in this case, or FrameTracker object will be leaked.
-        verify(mTracker).removeObservers();
+        verify(tracker).removeObservers();
 
         // Should never trigger Perfetto since it is a cancel.
-        verify(mTracker, never()).triggerPerfetto();
+        verify(tracker, never()).triggerPerfetto();
     }
 
     @Test
     public void testCancelWhenSessionNeverBegun() {
-        mTracker.cancel(FrameTracker.REASON_CANCEL_NORMAL);
-        verify(mTracker).removeObservers();
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
+        tracker.cancel(FrameTracker.REASON_CANCEL_NORMAL);
+        verify(tracker).removeObservers();
     }
 
     @Test
     public void testEndWhenSessionNeverBegun() {
-        mTracker.end(FrameTracker.REASON_END_NORMAL);
-        verify(mTracker).removeObservers();
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX, /* surfaceOnly= */ false);
+
+        tracker.end(FrameTracker.REASON_END_NORMAL);
+        verify(tracker).removeObservers();
     }
 
-    private void sendFirstWindowFrame(long durationMillis,
+    @Test
+    public void testSurfaceOnlyOtherFrameJanky() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_WALLPAPER_TRANSITION, CUJ_POSTFIX, /* surfaceOnly= */ true);
+
+        when(mChoreographer.getVsyncId()).thenReturn(100L);
+        tracker.begin();
+        verify(mSurfaceControlWrapper).addJankStatsListener(any(), any());
+
+        // First frame - not janky
+        sendFrame(tracker, JANK_NONE, 100L);
+        // normal frame - not janky
+        sendFrame(tracker, JANK_NONE, 101L);
+        // a janky frame
+        sendFrame(tracker, JANK_APP_DEADLINE_MISSED, 102L);
+
+        when(mChoreographer.getVsyncId()).thenReturn(102L);
+        tracker.end(FrameTracker.REASON_CANCEL_NORMAL);
+
+        // an extra frame to trigger finish
+        sendFrame(tracker, JANK_NONE, 103L);
+
+        verify(mSurfaceControlWrapper).removeJankStatsListener(any());
+        verify(tracker).triggerPerfetto();
+    }
+
+    @Test
+    public void testSurfaceOnlyFirstFrameJanky() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_WALLPAPER_TRANSITION, CUJ_POSTFIX, /* surfaceOnly= */ true);
+
+        when(mChoreographer.getVsyncId()).thenReturn(100L);
+        tracker.begin();
+        verify(mSurfaceControlWrapper).addJankStatsListener(any(), any());
+
+        // First frame - janky
+        sendFrame(tracker, JANK_APP_DEADLINE_MISSED, 100L);
+        // normal frame - not janky
+        sendFrame(tracker, JANK_NONE, 101L);
+        // normal frame - not janky
+        sendFrame(tracker, JANK_NONE, 102L);
+
+        when(mChoreographer.getVsyncId()).thenReturn(102L);
+        tracker.end(FrameTracker.REASON_CANCEL_NORMAL);
+
+        // an extra frame to trigger finish
+        sendFrame(tracker, JANK_NONE, 103L);
+
+        verify(mSurfaceControlWrapper).removeJankStatsListener(any());
+        verify(tracker, never()).triggerPerfetto();
+    }
+
+    @Test
+    public void testSurfaceOnlyLastFrameJanky() {
+        FrameTracker tracker = spyFrameTracker(
+                CUJ_WALLPAPER_TRANSITION, CUJ_POSTFIX, /* surfaceOnly= */ true);
+
+        when(mChoreographer.getVsyncId()).thenReturn(100L);
+        tracker.begin();
+        verify(mSurfaceControlWrapper).addJankStatsListener(any(), any());
+
+        // First frame - not janky
+        sendFrame(tracker, JANK_NONE, 100L);
+        // normal frame - not janky
+        sendFrame(tracker, JANK_NONE, 101L);
+        // normal frame - not janky
+        sendFrame(tracker, JANK_NONE, 102L);
+
+        when(mChoreographer.getVsyncId()).thenReturn(102L);
+        tracker.end(FrameTracker.REASON_CANCEL_NORMAL);
+
+        // janky frame, should be ignored, trigger finish
+        sendFrame(tracker, JANK_APP_DEADLINE_MISSED, 103L);
+
+        verify(mSurfaceControlWrapper).removeJankStatsListener(any());
+        verify(tracker, never()).triggerPerfetto();
+    }
+
+    private void sendFirstWindowFrame(FrameTracker tracker, long durationMillis,
             @JankType int jankType, long vsyncId) {
-        sendFrame(durationMillis, jankType, vsyncId, true /* firstWindowFrame */);
+        sendFrame(tracker, durationMillis, jankType, vsyncId, /* firstWindowFrame= */ true);
     }
 
-    private void sendFrame(long durationMillis,
+    private void sendFrame(FrameTracker tracker, long durationMillis,
             @JankType int jankType, long vsyncId) {
-        sendFrame(durationMillis, jankType, vsyncId, false /* firstWindowFrame */);
+        sendFrame(tracker, durationMillis, jankType, vsyncId, /* firstWindowFrame= */ false);
+    }
+
+    /**
+     * Used for surface only test.
+     */
+    private void sendFrame(FrameTracker tracker, @JankType int jankType, long vsyncId) {
+        sendFrame(tracker, /* durationMillis= */ -1,
+                jankType, vsyncId, /* firstWindowFrame= */ false);
     }
 
-    private void sendFrame(long durationMillis,
+    private void sendFrame(FrameTracker tracker, long durationMillis,
             @JankType int jankType, long vsyncId, boolean firstWindowFrame) {
-        when(mWrapper.getTiming()).thenReturn(new long[] { 0, vsyncId });
-        doReturn(firstWindowFrame ? 1L : 0L).when(mWrapper)
-                .getMetric(FrameMetrics.FIRST_DRAW_FRAME);
-        doReturn(TimeUnit.MILLISECONDS.toNanos(durationMillis))
-                .when(mWrapper).getMetric(FrameMetrics.TOTAL_DURATION);
-        mTracker.onFrameMetricsAvailable(0);
+        if (!tracker.mSurfaceOnly) {
+            when(mWrapper.getTiming()).thenReturn(new long[]{0, vsyncId});
+            doReturn(firstWindowFrame ? 1L : 0L).when(mWrapper)
+                    .getMetric(FrameMetrics.FIRST_DRAW_FRAME);
+            doReturn(TimeUnit.MILLISECONDS.toNanos(durationMillis))
+                    .when(mWrapper).getMetric(FrameMetrics.TOTAL_DURATION);
+            tracker.onFrameMetricsAvailable(0);
+        }
         mListenerCapture.getValue().onJankDataAvailable(new JankData[] {
                 new JankData(vsyncId, jankType)
         });
diff --git a/core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java b/core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java
index c153b38d3f02..0d2d047b7f0b 100644
--- a/core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java
+++ b/core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java
@@ -16,8 +16,8 @@
 
 package com.android.internal.jank;
 
-import static com.android.internal.jank.FrameTracker.SurfaceControlWrapper;
-import static com.android.internal.jank.FrameTracker.ViewRootWrapper;
+import static com.android.internal.jank.FrameTracker.REASON_CANCEL_TIMEOUT;
+import static com.android.internal.jank.FrameTracker.REASON_END_NORMAL;
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE;
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_TO_STATSD_INTERACTION_TYPE;
 
@@ -25,17 +25,18 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
-import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.os.Handler;
 import android.os.HandlerThread;
-import android.os.Message;
+import android.os.SystemClock;
 import android.provider.DeviceConfig;
 import android.view.View;
 import android.view.ViewAttachTestActivity;
@@ -43,8 +44,12 @@ import android.view.ViewAttachTestActivity;
 import androidx.test.filters.SmallTest;
 import androidx.test.rule.ActivityTestRule;
 
+import com.android.internal.jank.FrameTracker.ChoreographerWrapper;
 import com.android.internal.jank.FrameTracker.FrameMetricsWrapper;
+import com.android.internal.jank.FrameTracker.SurfaceControlWrapper;
 import com.android.internal.jank.FrameTracker.ThreadedRendererWrapper;
+import com.android.internal.jank.FrameTracker.ViewRootWrapper;
+import com.android.internal.jank.InteractionJankMonitor.Configuration;
 import com.android.internal.jank.InteractionJankMonitor.Session;
 
 import org.junit.Before;
@@ -80,33 +85,23 @@ public class InteractionJankMonitorTest {
 
         Handler handler = spy(new Handler(mActivity.getMainLooper()));
         doReturn(true).when(handler).sendMessageAtTime(any(), anyLong());
-        mWorker = spy(new HandlerThread("Interaction-jank-monitor-test"));
-        doNothing().when(mWorker).start();
+        mWorker = mock(HandlerThread.class);
         doReturn(handler).when(mWorker).getThreadHandler();
     }
 
     @Test
     public void testBeginEnd() {
-        // Should return false if the view is not attached.
-        InteractionJankMonitor monitor = spy(new InteractionJankMonitor(mWorker));
-        verify(mWorker).start();
-
-        Session session = new Session(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX);
-        FrameTracker tracker = spy(new FrameTracker(session, mWorker.getThreadHandler(),
-                new ThreadedRendererWrapper(mView.getThreadedRenderer()),
-                new ViewRootWrapper(mView.getViewRootImpl()), new SurfaceControlWrapper(),
-                mock(FrameTracker.ChoreographerWrapper.class),
-                new FrameMetricsWrapper(), /*traceThresholdMissedFrames=*/ 1,
-                /*traceThresholdFrameTimeMillis=*/ -1, null));
+        InteractionJankMonitor monitor = createMockedInteractionJankMonitor();
+        FrameTracker tracker = createMockedFrameTracker(null);
         doReturn(tracker).when(monitor).createFrameTracker(any(), any());
-        doNothing().when(tracker).triggerPerfetto();
-        doNothing().when(tracker).postTraceStartMarker();
+        doNothing().when(tracker).begin();
+        doReturn(true).when(tracker).end(anyInt());
 
         // Simulate a trace session and see if begin / end are invoked.
-        assertThat(monitor.begin(mView, session.getCuj())).isTrue();
+        assertThat(monitor.begin(mView, CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE)).isTrue();
         verify(tracker).begin();
-        assertThat(monitor.end(session.getCuj())).isTrue();
-        verify(tracker).end(FrameTracker.REASON_END_NORMAL);
+        assertThat(monitor.end(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE)).isTrue();
+        verify(tracker).end(REASON_END_NORMAL);
     }
 
     @Test
@@ -135,31 +130,23 @@ public class InteractionJankMonitorTest {
     }
 
     @Test
-    public void testBeginCancel() {
-        InteractionJankMonitor monitor = spy(new InteractionJankMonitor(mWorker));
-
-        ArgumentCaptor<Message> captor = ArgumentCaptor.forClass(Message.class);
-
-        Session session = new Session(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX);
-        FrameTracker tracker = spy(new FrameTracker(session, mWorker.getThreadHandler(),
-                new ThreadedRendererWrapper(mView.getThreadedRenderer()),
-                new ViewRootWrapper(mView.getViewRootImpl()), new SurfaceControlWrapper(),
-                mock(FrameTracker.ChoreographerWrapper.class),
-                new FrameMetricsWrapper(), /*traceThresholdMissedFrames=*/ 1,
-                /*traceThresholdFrameTimeMillis=*/ -1, null));
+    public void testBeginTimeout() {
+        ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+        InteractionJankMonitor monitor = createMockedInteractionJankMonitor();
+        FrameTracker tracker = createMockedFrameTracker(null);
         doReturn(tracker).when(monitor).createFrameTracker(any(), any());
-        doNothing().when(tracker).triggerPerfetto();
-        doNothing().when(tracker).postTraceStartMarker();
+        doNothing().when(tracker).begin();
+        doReturn(true).when(tracker).cancel(anyInt());
 
-        assertThat(monitor.begin(mView, session.getCuj())).isTrue();
+        assertThat(monitor.begin(mView, CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE)).isTrue();
         verify(tracker).begin();
-        verify(mWorker.getThreadHandler(), atLeastOnce()).sendMessageAtTime(captor.capture(),
-                anyLong());
-        Runnable runnable = captor.getValue().getCallback();
+        verify(monitor).scheduleTimeoutAction(anyInt(), anyLong(), captor.capture());
+        Runnable runnable = captor.getValue();
         assertThat(runnable).isNotNull();
         mWorker.getThreadHandler().removeCallbacks(runnable);
         runnable.run();
-        verify(tracker).cancel(FrameTracker.REASON_CANCEL_NORMAL);
+        verify(monitor).cancel(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, REASON_CANCEL_TIMEOUT);
+        verify(tracker).cancel(REASON_CANCEL_TIMEOUT);
     }
 
     @Test
@@ -185,4 +172,43 @@ public class InteractionJankMonitorTest {
                     .isTrue();
         }
     }
+
+    private InteractionJankMonitor createMockedInteractionJankMonitor() {
+        InteractionJankMonitor monitor = spy(new InteractionJankMonitor(mWorker));
+        doReturn(true).when(monitor).shouldMonitor(anyInt());
+        doNothing().when(monitor).notifyEvents(any(), any(), any());
+        return monitor;
+    }
+
+    private FrameTracker createMockedFrameTracker(FrameTracker.FrameTrackerListener listener) {
+        Session session = spy(new Session(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX));
+        doReturn(false).when(session).logToStatsd();
+
+        ThreadedRendererWrapper threadedRenderer = mock(ThreadedRendererWrapper.class);
+        doNothing().when(threadedRenderer).addObserver(any());
+        doNothing().when(threadedRenderer).removeObserver(any());
+
+        ViewRootWrapper viewRoot = spy(new ViewRootWrapper(mView.getViewRootImpl()));
+        doNothing().when(viewRoot).addSurfaceChangedCallback(any());
+
+        SurfaceControlWrapper surfaceControl = mock(SurfaceControlWrapper.class);
+        doNothing().when(surfaceControl).addJankStatsListener(any(), any());
+        doNothing().when(surfaceControl).removeJankStatsListener(any());
+
+        final ChoreographerWrapper choreographer = mock(ChoreographerWrapper.class);
+        doReturn(SystemClock.elapsedRealtime()).when(choreographer).getVsyncId();
+
+        Configuration configuration = mock(Configuration.class);
+        when(configuration.isSurfaceOnly()).thenReturn(false);
+
+        FrameTracker tracker = spy(new FrameTracker(session, mWorker.getThreadHandler(),
+                threadedRenderer, viewRoot, surfaceControl, choreographer,
+                new FrameMetricsWrapper(), /* traceThresholdMissedFrames= */ 1,
+                /* traceThresholdFrameTimeMillis= */ -1, listener, configuration));
+
+        doNothing().when(tracker).postTraceStartMarker();
+        doNothing().when(tracker).triggerPerfetto();
+
+        return tracker;
+    }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
index 79f7a5c9df18..130f552f6e3a 100644
--- a/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
@@ -16,6 +16,8 @@
 
 package com.android.internal.os;
 
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_AMBIENT;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import android.os.BatteryConsumer;
@@ -36,26 +38,28 @@ public class AmbientDisplayPowerCalculatorTest {
 
     @Rule
     public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
-            .setAveragePower(PowerProfile.POWER_AMBIENT_DISPLAY, 10.0);
+            .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 0, 10.0)
+            .setNumDisplays(1);
 
     @Test
     public void testMeasuredEnergyBasedModel() {
         mStatsRule.initMeasuredEnergyStatsLocked();
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
-        stats.updateDisplayMeasuredEnergyStatsLocked(300_000_000, Display.STATE_ON, 0);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{300_000_000},
+                new int[]{Display.STATE_ON}, 0);
 
-        stats.noteScreenStateLocked(Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
+        stats.noteScreenStateLocked(0, Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
                 30 * MINUTE_IN_MS);
 
-        stats.updateDisplayMeasuredEnergyStatsLocked(200_000_000, Display.STATE_DOZE,
-                30 * MINUTE_IN_MS);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{200_000_000},
+                new int[]{Display.STATE_DOZE}, 30 * MINUTE_IN_MS);
 
-        stats.noteScreenStateLocked(Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
+        stats.noteScreenStateLocked(0, Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
                 120 * MINUTE_IN_MS);
 
-        stats.updateDisplayMeasuredEnergyStatsLocked(100_000_000, Display.STATE_OFF,
-                120 * MINUTE_IN_MS);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{100_000_000},
+                new int[]{Display.STATE_OFF}, 120 * MINUTE_IN_MS);
 
         AmbientDisplayPowerCalculator calculator =
                 new AmbientDisplayPowerCalculator(mStatsRule.getPowerProfile());
@@ -73,12 +77,73 @@ public class AmbientDisplayPowerCalculatorTest {
     }
 
     @Test
+    public void testMeasuredEnergyBasedModel_multiDisplay() {
+        mStatsRule.initMeasuredEnergyStatsLocked()
+                .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 1, 20.0)
+                .setNumDisplays(2);
+        BatteryStatsImpl stats = mStatsRule.getBatteryStats();
+
+
+        final int[] screenStates = new int[] {Display.STATE_OFF, Display.STATE_OFF};
+
+        stats.noteScreenStateLocked(0, screenStates[0], 0, 0, 0);
+        stats.noteScreenStateLocked(1, screenStates[1], 0, 0, 0);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{300, 400}, screenStates, 0);
+
+        // Switch display0 to doze
+        screenStates[0] = Display.STATE_DOZE;
+        stats.noteScreenStateLocked(0, screenStates[0], 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
+                30 * MINUTE_IN_MS);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{200, 300},
+                screenStates, 30 * MINUTE_IN_MS);
+
+        // Switch display1 to doze
+        screenStates[1] = Display.STATE_DOZE;
+        stats.noteScreenStateLocked(1, Display.STATE_DOZE, 90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS,
+                90 * MINUTE_IN_MS);
+        // 100,000,000 uC should be attributed to display 0 doze here.
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{100_000_000, 700_000_000},
+                screenStates, 90 * MINUTE_IN_MS);
+
+        // Switch display0 to off
+        screenStates[0] = Display.STATE_OFF;
+        stats.noteScreenStateLocked(0, screenStates[0], 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
+                120 * MINUTE_IN_MS);
+        // 40,000,000 and 70,000,000 uC should be attributed to display 0 and 1 doze here.
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{40_000_000, 70_000_000},
+                screenStates, 120 * MINUTE_IN_MS);
+
+        // Switch display1 to off
+        screenStates[1] = Display.STATE_OFF;
+        stats.noteScreenStateLocked(1, screenStates[1], 150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS,
+                150 * MINUTE_IN_MS);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{100, 90_000_000}, screenStates,
+                150 * MINUTE_IN_MS);
+        // 90,000,000 uC should be attributed to display 1 doze here.
+
+        AmbientDisplayPowerCalculator calculator =
+                new AmbientDisplayPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(calculator);
+
+        BatteryConsumer consumer = mStatsRule.getDeviceBatteryConsumer();
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isEqualTo(120 * MINUTE_IN_MS);
+        // 100,000,000 + 40,000,000 + 70,000,000 + 90,000,000 uC / 1000 (micro-/milli-) / 3600
+        // (seconds/hour) = 27.777778 mAh
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isWithin(PRECISION).of(83.33333);
+        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+    }
+
+    @Test
     public void testPowerProfileBasedModel() {
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
-        stats.noteScreenStateLocked(Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
+        stats.noteScreenStateLocked(0, Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
                 30 * MINUTE_IN_MS);
-        stats.noteScreenStateLocked(Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
+        stats.noteScreenStateLocked(0, Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
                 120 * MINUTE_IN_MS);
 
         AmbientDisplayPowerCalculator calculator =
@@ -94,4 +159,36 @@ public class AmbientDisplayPowerCalculatorTest {
         assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
     }
+
+    @Test
+    public void testPowerProfileBasedModel_multiDisplay() {
+        mStatsRule.setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 1, 20.0)
+                .setNumDisplays(2);
+
+        BatteryStatsImpl stats = mStatsRule.getBatteryStats();
+
+        stats.noteScreenStateLocked(1, Display.STATE_OFF, 0, 0, 0);
+        stats.noteScreenStateLocked(0, Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
+                30 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, Display.STATE_DOZE, 90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS,
+                90 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
+                120 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, Display.STATE_OFF, 150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS,
+                150 * MINUTE_IN_MS);
+
+        AmbientDisplayPowerCalculator calculator =
+                new AmbientDisplayPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(BatteryUsageStatsRule.POWER_PROFILE_MODEL_ONLY, calculator);
+
+        BatteryConsumer consumer = mStatsRule.getDeviceBatteryConsumer();
+        // Duration should only be the union of
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isEqualTo(120 * MINUTE_IN_MS);
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isWithin(PRECISION).of(35.0);
+        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+    }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
index d4799a8f5fd3..3e2885a74287 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
@@ -16,9 +16,13 @@
 
 package com.android.internal.os;
 
+import static android.os.BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS;
 import static android.os.BatteryStats.STATS_SINCE_CHARGED;
 import static android.os.BatteryStats.WAKE_TYPE_PARTIAL;
 
+import static com.android.internal.os.BatteryStatsImpl.ExternalStatsSync.UPDATE_CPU;
+import static com.android.internal.os.BatteryStatsImpl.ExternalStatsSync.UPDATE_DISPLAY;
+
 import android.app.ActivityManager;
 import android.os.BatteryStats;
 import android.os.BatteryStats.HistoryItem;
@@ -37,8 +41,10 @@ import com.android.internal.power.MeasuredEnergyStats;
 
 import junit.framework.TestCase;
 
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.IntConsumer;
 
 /**
  * Test various BatteryStatsImpl noteStart methods.
@@ -317,18 +323,130 @@ public class BatteryStatsNoteTest extends TestCase {
     public void testNoteScreenStateLocked() throws Exception {
         final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
         MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+        bi.initMeasuredEnergyStats(new String[]{"FOO", "BAR"});
 
         bi.updateTimeBasesLocked(true, Display.STATE_ON, 0, 0);
-        bi.noteScreenStateLocked(Display.STATE_ON);
-        bi.noteScreenStateLocked(Display.STATE_DOZE);
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
+
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_DOZE, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_OFF, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE_SUSPEND);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_DOZE_SUSPEND, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        // STATE_VR note should map to STATE_ON.
+        bi.noteScreenStateLocked(0, Display.STATE_VR);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        // STATE_ON_SUSPEND note should map to STATE_ON.
+        bi.noteScreenStateLocked(0, Display.STATE_ON_SUSPEND);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        // Transition from ON to ON state should not cause an External Sync
+        assertEquals(0, bi.getAndClearExternalStatsSyncFlags());
+    }
+
+    /**
+     * Test BatteryStatsImpl.noteScreenStateLocked sets timebases and screen states correctly for
+     * multi display devices
+     */
+    @SmallTest
+    public void testNoteScreenStateLocked_multiDisplay() throws Exception {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+        bi.setDisplayCountLocked(2);
+        bi.initMeasuredEnergyStats(new String[]{"FOO", "BAR"});
+
+        bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
+        bi.noteScreenStateLocked(1, Display.STATE_OFF);
+
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_DOZE, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_OFF, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE_SUSPEND);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_DOZE_SUSPEND, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        // STATE_VR note should map to STATE_ON.
+        bi.noteScreenStateLocked(0, Display.STATE_VR);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        // STATE_ON_SUSPEND note should map to STATE_ON.
+        bi.noteScreenStateLocked(0, Display.STATE_ON_SUSPEND);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        // Transition from ON to ON state should not cause an External Sync
+        assertEquals(0, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(1, Display.STATE_DOZE);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        // Should remain STATE_ON since display0 is still on.
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        // Overall screen state did not change, so no need to sync CPU stats.
+        assertEquals(UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE);
         assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
-        assertEquals(bi.getScreenState(), Display.STATE_DOZE);
-        bi.noteScreenStateLocked(Display.STATE_ON);
+        assertEquals(Display.STATE_DOZE, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
         assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
-        assertEquals(bi.getScreenState(), Display.STATE_ON);
-        bi.noteScreenStateLocked(Display.STATE_OFF);
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
+        assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_DOZE, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE_SUSPEND);
         assertTrue(bi.getOnBatteryScreenOffTimeBase().isRunning());
-        assertEquals(bi.getScreenState(), Display.STATE_OFF);
+        assertEquals(Display.STATE_DOZE, bi.getScreenState());
+        // Overall screen state did not change, so no need to sync CPU stats.
+        assertEquals(UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_VR);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        assertEquals(UPDATE_CPU | UPDATE_DISPLAY, bi.getAndClearExternalStatsSyncFlags());
+
+        bi.noteScreenStateLocked(0, Display.STATE_ON_SUSPEND);
+        assertFalse(bi.getOnBatteryScreenOffTimeBase().isRunning());
+        assertEquals(Display.STATE_ON, bi.getScreenState());
+        assertEquals(0, bi.getAndClearExternalStatsSyncFlags());
     }
 
     /*
@@ -352,32 +470,317 @@ public class BatteryStatsNoteTest extends TestCase {
         bi.updateTimeBasesLocked(true, Display.STATE_UNKNOWN, 100_000, 100_000);
         // Turn on display at 200us
         clocks.realtime = clocks.uptime = 200;
-        bi.noteScreenStateLocked(Display.STATE_ON);
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
         assertEquals(150_000, bi.computeBatteryRealtime(250_000, STATS_SINCE_CHARGED));
         assertEquals(100_000, bi.computeBatteryScreenOffRealtime(250_000, STATS_SINCE_CHARGED));
         assertEquals(50_000, bi.getScreenOnTime(250_000, STATS_SINCE_CHARGED));
         assertEquals(0, bi.getScreenDozeTime(250_000, STATS_SINCE_CHARGED));
+        assertEquals(50_000, bi.getDisplayScreenOnTime(0, 250_000));
+        assertEquals(0, bi.getDisplayScreenDozeTime(0, 250_000));
 
         clocks.realtime = clocks.uptime = 310;
-        bi.noteScreenStateLocked(Display.STATE_OFF);
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
         assertEquals(250_000, bi.computeBatteryRealtime(350_000, STATS_SINCE_CHARGED));
         assertEquals(140_000, bi.computeBatteryScreenOffRealtime(350_000, STATS_SINCE_CHARGED));
         assertEquals(110_000, bi.getScreenOnTime(350_000, STATS_SINCE_CHARGED));
         assertEquals(0, bi.getScreenDozeTime(350_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getDisplayScreenOnTime(0, 350_000));
+        assertEquals(0, bi.getDisplayScreenDozeTime(0, 350_000));
 
         clocks.realtime = clocks.uptime = 400;
-        bi.noteScreenStateLocked(Display.STATE_DOZE);
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE);
         assertEquals(400_000, bi.computeBatteryRealtime(500_000, STATS_SINCE_CHARGED));
         assertEquals(290_000, bi.computeBatteryScreenOffRealtime(500_000, STATS_SINCE_CHARGED));
         assertEquals(110_000, bi.getScreenOnTime(500_000, STATS_SINCE_CHARGED));
         assertEquals(100_000, bi.getScreenDozeTime(500_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getDisplayScreenOnTime(0, 500_000));
+        assertEquals(100_000, bi.getDisplayScreenDozeTime(0, 500_000));
 
         clocks.realtime = clocks.uptime = 1000;
-        bi.noteScreenStateLocked(Display.STATE_OFF);
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
         assertEquals(1400_000, bi.computeBatteryRealtime(1500_000, STATS_SINCE_CHARGED));
         assertEquals(1290_000, bi.computeBatteryScreenOffRealtime(1500_000, STATS_SINCE_CHARGED));
         assertEquals(110_000, bi.getScreenOnTime(1500_000, STATS_SINCE_CHARGED));
         assertEquals(600_000, bi.getScreenDozeTime(1500_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getDisplayScreenOnTime(0, 1500_000));
+        assertEquals(600_000, bi.getDisplayScreenDozeTime(0, 1500_000));
+    }
+
+    /*
+     * Test BatteryStatsImpl.noteScreenStateLocked updates timers correctly for multi display
+     * devices.
+     */
+    @SmallTest
+    public void testNoteScreenStateTimersLocked_multiDisplay() throws Exception {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+        bi.setDisplayCountLocked(2);
+
+        clocks.realtime = clocks.uptime = 100;
+        // Device startup, setOnBatteryLocked calls updateTimebases
+        bi.updateTimeBasesLocked(true, Display.STATE_UNKNOWN, 100_000, 100_000);
+        // Turn on display at 200us
+        clocks.realtime = clocks.uptime = 200;
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
+        bi.noteScreenStateLocked(1, Display.STATE_OFF);
+        assertEquals(150_000, bi.computeBatteryRealtime(250_000, STATS_SINCE_CHARGED));
+        assertEquals(100_000, bi.computeBatteryScreenOffRealtime(250_000, STATS_SINCE_CHARGED));
+        assertEquals(50_000, bi.getScreenOnTime(250_000, STATS_SINCE_CHARGED));
+        assertEquals(0, bi.getScreenDozeTime(250_000, STATS_SINCE_CHARGED));
+        assertEquals(50_000, bi.getDisplayScreenOnTime(0, 250_000));
+        assertEquals(0, bi.getDisplayScreenDozeTime(0, 250_000));
+        assertEquals(0, bi.getDisplayScreenOnTime(1, 250_000));
+        assertEquals(0, bi.getDisplayScreenDozeTime(1, 250_000));
+
+        clocks.realtime = clocks.uptime = 310;
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
+        assertEquals(250_000, bi.computeBatteryRealtime(350_000, STATS_SINCE_CHARGED));
+        assertEquals(140_000, bi.computeBatteryScreenOffRealtime(350_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getScreenOnTime(350_000, STATS_SINCE_CHARGED));
+        assertEquals(0, bi.getScreenDozeTime(350_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getDisplayScreenOnTime(0, 350_000));
+        assertEquals(0, bi.getDisplayScreenDozeTime(0, 350_000));
+        assertEquals(0, bi.getDisplayScreenOnTime(1, 350_000));
+        assertEquals(0, bi.getDisplayScreenDozeTime(1, 350_000));
+
+        clocks.realtime = clocks.uptime = 400;
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE);
+        assertEquals(400_000, bi.computeBatteryRealtime(500_000, STATS_SINCE_CHARGED));
+        assertEquals(290_000, bi.computeBatteryScreenOffRealtime(500_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getScreenOnTime(500_000, STATS_SINCE_CHARGED));
+        assertEquals(100_000, bi.getScreenDozeTime(500_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getDisplayScreenOnTime(0, 500_000));
+        assertEquals(100_000, bi.getDisplayScreenDozeTime(0, 500_000));
+        assertEquals(0, bi.getDisplayScreenOnTime(1, 500_000));
+        assertEquals(0, bi.getDisplayScreenDozeTime(1, 500_000));
+
+        clocks.realtime = clocks.uptime = 1000;
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
+        assertEquals(1000_000, bi.computeBatteryRealtime(1100_000, STATS_SINCE_CHARGED));
+        assertEquals(890_000, bi.computeBatteryScreenOffRealtime(1100_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getScreenOnTime(1100_000, STATS_SINCE_CHARGED));
+        assertEquals(600_000, bi.getScreenDozeTime(1100_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getDisplayScreenOnTime(0, 1100_000));
+        assertEquals(600_000, bi.getDisplayScreenDozeTime(0, 1100_000));
+        assertEquals(0, bi.getDisplayScreenOnTime(1, 1100_000));
+        assertEquals(0, bi.getDisplayScreenDozeTime(1, 1100_000));
+
+        clocks.realtime = clocks.uptime = 1200;
+        // Change state of second display to doze
+        bi.noteScreenStateLocked(1, Display.STATE_DOZE);
+        assertEquals(1150_000, bi.computeBatteryRealtime(1250_000, STATS_SINCE_CHARGED));
+        assertEquals(1040_000, bi.computeBatteryScreenOffRealtime(1250_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getScreenOnTime(1250_000, STATS_SINCE_CHARGED));
+        assertEquals(650_000, bi.getScreenDozeTime(1250_000, STATS_SINCE_CHARGED));
+        assertEquals(110_000, bi.getDisplayScreenOnTime(0, 1250_000));
+        assertEquals(600_000, bi.getDisplayScreenDozeTime(0, 1250_000));
+        assertEquals(0, bi.getDisplayScreenOnTime(1, 1250_000));
+        assertEquals(50_000, bi.getDisplayScreenDozeTime(1, 1250_000));
+
+        clocks.realtime = clocks.uptime = 1310;
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
+        assertEquals(1250_000, bi.computeBatteryRealtime(1350_000, STATS_SINCE_CHARGED));
+        assertEquals(1100_000, bi.computeBatteryScreenOffRealtime(1350_000, STATS_SINCE_CHARGED));
+        assertEquals(150_000, bi.getScreenOnTime(1350_000, STATS_SINCE_CHARGED));
+        assertEquals(710_000, bi.getScreenDozeTime(1350_000, STATS_SINCE_CHARGED));
+        assertEquals(150_000, bi.getDisplayScreenOnTime(0, 1350_000));
+        assertEquals(600_000, bi.getDisplayScreenDozeTime(0, 1350_000));
+        assertEquals(0, bi.getDisplayScreenOnTime(1, 1350_000));
+        assertEquals(150_000, bi.getDisplayScreenDozeTime(1, 1350_000));
+
+        clocks.realtime = clocks.uptime = 1400;
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE);
+        assertEquals(1400_000, bi.computeBatteryRealtime(1500_000, STATS_SINCE_CHARGED));
+        assertEquals(1200_000, bi.computeBatteryScreenOffRealtime(1500_000, STATS_SINCE_CHARGED));
+        assertEquals(200_000, bi.getScreenOnTime(1500_000, STATS_SINCE_CHARGED));
+        assertEquals(810_000, bi.getScreenDozeTime(1500_000, STATS_SINCE_CHARGED));
+        assertEquals(200_000, bi.getDisplayScreenOnTime(0, 1500_000));
+        assertEquals(700_000, bi.getDisplayScreenDozeTime(0, 1500_000));
+        assertEquals(0, bi.getDisplayScreenOnTime(1, 1500_000));
+        assertEquals(300_000, bi.getDisplayScreenDozeTime(1, 1500_000));
+
+        clocks.realtime = clocks.uptime = 2000;
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
+        assertEquals(2000_000, bi.computeBatteryRealtime(2100_000, STATS_SINCE_CHARGED));
+        assertEquals(1800_000, bi.computeBatteryScreenOffRealtime(2100_000, STATS_SINCE_CHARGED));
+        assertEquals(200_000, bi.getScreenOnTime(2100_000, STATS_SINCE_CHARGED));
+        assertEquals(1410_000, bi.getScreenDozeTime(2100_000, STATS_SINCE_CHARGED));
+        assertEquals(200_000, bi.getDisplayScreenOnTime(0, 2100_000));
+        assertEquals(1200_000, bi.getDisplayScreenDozeTime(0, 2100_000));
+        assertEquals(0, bi.getDisplayScreenOnTime(1, 2100_000));
+        assertEquals(900_000, bi.getDisplayScreenDozeTime(1, 2100_000));
+
+
+        clocks.realtime = clocks.uptime = 2200;
+        // Change state of second display to on
+        bi.noteScreenStateLocked(1, Display.STATE_ON);
+        assertEquals(2150_000, bi.computeBatteryRealtime(2250_000, STATS_SINCE_CHARGED));
+        assertEquals(1900_000, bi.computeBatteryScreenOffRealtime(2250_000, STATS_SINCE_CHARGED));
+        assertEquals(250_000, bi.getScreenOnTime(2250_000, STATS_SINCE_CHARGED));
+        assertEquals(1510_000, bi.getScreenDozeTime(2250_000, STATS_SINCE_CHARGED));
+        assertEquals(200_000, bi.getDisplayScreenOnTime(0, 2250_000));
+        assertEquals(1200_000, bi.getDisplayScreenDozeTime(0, 2250_000));
+        assertEquals(50_000, bi.getDisplayScreenOnTime(1, 2250_000));
+        assertEquals(1000_000, bi.getDisplayScreenDozeTime(1, 2250_000));
+
+        clocks.realtime = clocks.uptime = 2310;
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
+        assertEquals(2250_000, bi.computeBatteryRealtime(2350_000, STATS_SINCE_CHARGED));
+        assertEquals(1900_000, bi.computeBatteryScreenOffRealtime(2350_000, STATS_SINCE_CHARGED));
+        assertEquals(350_000, bi.getScreenOnTime(2350_000, STATS_SINCE_CHARGED));
+        assertEquals(1510_000, bi.getScreenDozeTime(2350_000, STATS_SINCE_CHARGED));
+        assertEquals(240_000, bi.getDisplayScreenOnTime(0, 2350_000));
+        assertEquals(1200_000, bi.getDisplayScreenDozeTime(0, 2350_000));
+        assertEquals(150_000, bi.getDisplayScreenOnTime(1, 2350_000));
+        assertEquals(1000_000, bi.getDisplayScreenDozeTime(1, 2350_000));
+
+        clocks.realtime = clocks.uptime = 2400;
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE);
+        assertEquals(2400_000, bi.computeBatteryRealtime(2500_000, STATS_SINCE_CHARGED));
+        assertEquals(1900_000, bi.computeBatteryScreenOffRealtime(2500_000, STATS_SINCE_CHARGED));
+        assertEquals(500_000, bi.getScreenOnTime(2500_000, STATS_SINCE_CHARGED));
+        assertEquals(1510_000, bi.getScreenDozeTime(2500_000, STATS_SINCE_CHARGED));
+        assertEquals(290_000, bi.getDisplayScreenOnTime(0, 2500_000));
+        assertEquals(1300_000, bi.getDisplayScreenDozeTime(0, 2500_000));
+        assertEquals(300_000, bi.getDisplayScreenOnTime(1, 2500_000));
+        assertEquals(1000_000, bi.getDisplayScreenDozeTime(1, 2500_000));
+
+        clocks.realtime = clocks.uptime = 3000;
+        bi.noteScreenStateLocked(0, Display.STATE_OFF);
+        assertEquals(3000_000, bi.computeBatteryRealtime(3100_000, STATS_SINCE_CHARGED));
+        assertEquals(1900_000, bi.computeBatteryScreenOffRealtime(3100_000, STATS_SINCE_CHARGED));
+        assertEquals(1100_000, bi.getScreenOnTime(3100_000, STATS_SINCE_CHARGED));
+        assertEquals(1510_000, bi.getScreenDozeTime(3100_000, STATS_SINCE_CHARGED));
+        assertEquals(290_000, bi.getDisplayScreenOnTime(0, 3100_000));
+        assertEquals(1800_000, bi.getDisplayScreenDozeTime(0, 3100_000));
+        assertEquals(900_000, bi.getDisplayScreenOnTime(1, 3100_000));
+        assertEquals(1000_000, bi.getDisplayScreenDozeTime(1, 3100_000));
+    }
+
+
+    /**
+     * Test BatteryStatsImpl.noteScreenBrightnessLocked updates timers correctly.
+     */
+    @SmallTest
+    public void testScreenBrightnessLocked_multiDisplay() throws Exception {
+        final MockClocks clocks = new MockClocks(); // holds realtime and uptime in ms
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+
+        final int numDisplay = 2;
+        bi.setDisplayCountLocked(numDisplay);
+
+
+        final long[] overallExpected = new long[NUM_SCREEN_BRIGHTNESS_BINS];
+        final long[][] perDisplayExpected = new long[numDisplay][NUM_SCREEN_BRIGHTNESS_BINS];
+        class Bookkeeper {
+            public long currentTimeMs = 100;
+            public int overallActiveBin = -1;
+            public int[] perDisplayActiveBin = new int[numDisplay];
+        }
+        final Bookkeeper bk = new Bookkeeper();
+        Arrays.fill(bk.perDisplayActiveBin, -1);
+
+        IntConsumer incrementTime = inc -> {
+            bk.currentTimeMs += inc;
+            if (bk.overallActiveBin >= 0) {
+                overallExpected[bk.overallActiveBin] += inc;
+            }
+            for (int i = 0; i < numDisplay; i++) {
+                final int bin = bk.perDisplayActiveBin[i];
+                if (bin >= 0) {
+                    perDisplayExpected[i][bin] += inc;
+                }
+            }
+            clocks.realtime = clocks.uptime = bk.currentTimeMs;
+        };
+
+        bi.updateTimeBasesLocked(true, Display.STATE_ON, 0, 0);
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
+        bi.noteScreenStateLocked(1, Display.STATE_ON);
+
+        incrementTime.accept(100);
+        bi.noteScreenBrightnessLocked(0, 25);
+        bi.noteScreenBrightnessLocked(1, 25);
+        // floor(25/256*5) = bin 0
+        bk.overallActiveBin = 0;
+        bk.perDisplayActiveBin[0] = 0;
+        bk.perDisplayActiveBin[1] = 0;
+
+        incrementTime.accept(50);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
+
+        incrementTime.accept(13);
+        bi.noteScreenBrightnessLocked(0, 100);
+        // floor(25/256*5) = bin 1
+        bk.overallActiveBin = 1;
+        bk.perDisplayActiveBin[0] = 1;
+
+        incrementTime.accept(44);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
+
+        incrementTime.accept(22);
+        bi.noteScreenBrightnessLocked(1, 200);
+        // floor(200/256*5) = bin 3
+        bk.overallActiveBin = 3;
+        bk.perDisplayActiveBin[1] = 3;
+
+        incrementTime.accept(33);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
+
+        incrementTime.accept(77);
+        bi.noteScreenBrightnessLocked(0, 150);
+        // floor(150/256*5) = bin 2
+        // Overall active bin should not change
+        bk.perDisplayActiveBin[0] = 2;
+
+        incrementTime.accept(88);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
+
+        incrementTime.accept(11);
+        bi.noteScreenStateLocked(1, Display.STATE_OFF);
+        // Display 1 should timers should stop incrementing
+        // Overall active bin should fallback to display 0's bin
+        bk.overallActiveBin = 2;
+        bk.perDisplayActiveBin[1] = -1;
+
+        incrementTime.accept(99);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
+
+        incrementTime.accept(200);
+        bi.noteScreenBrightnessLocked(0, 255);
+        // floor(150/256*5) = bin 4
+        bk.overallActiveBin = 4;
+        bk.perDisplayActiveBin[0] = 4;
+
+        incrementTime.accept(300);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
+
+        incrementTime.accept(200);
+        bi.noteScreenStateLocked(0, Display.STATE_DOZE);
+        // No displays are on. No brightness timers should be active.
+        bk.overallActiveBin = -1;
+        bk.perDisplayActiveBin[0] = -1;
+
+        incrementTime.accept(300);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
+
+        incrementTime.accept(400);
+        bi.noteScreenStateLocked(1, Display.STATE_ON);
+        // Display 1 turned back on.
+        bk.overallActiveBin = 3;
+        bk.perDisplayActiveBin[1] = 3;
+
+        incrementTime.accept(500);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
+
+        incrementTime.accept(600);
+        bi.noteScreenStateLocked(0, Display.STATE_ON);
+        // Display 0 turned back on.
+        bk.overallActiveBin = 4;
+        bk.perDisplayActiveBin[0] = 4;
+
+        incrementTime.accept(700);
+        checkScreenBrightnesses(overallExpected, perDisplayExpected, bi, bk.currentTimeMs);
     }
 
     @SmallTest
@@ -595,7 +998,7 @@ public class BatteryStatsNoteTest extends TestCase {
         bi.initMeasuredEnergyStats(new String[]{"FOO", "BAR"});
 
         clocks.realtime = 0;
-        int screen = Display.STATE_OFF;
+        int[] screen = new int[]{Display.STATE_OFF};
         boolean battery = false;
 
         final int uid1 = 10500;
@@ -605,35 +1008,35 @@ public class BatteryStatsNoteTest extends TestCase {
         long globalDoze = 0;
 
         // Case A: uid1 off, uid2 off, battery off, screen off
-        bi.updateTimeBasesLocked(battery, screen, clocks.realtime*1000, 0);
+        bi.updateTimeBasesLocked(battery, screen[0], clocks.realtime * 1000, 0);
         bi.setOnBatteryInternal(battery);
-        bi.updateDisplayMeasuredEnergyStatsLocked(500_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{500_000}, screen, clocks.realtime);
         checkMeasuredCharge("A", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case B: uid1 off, uid2 off, battery ON,  screen off
         clocks.realtime += 17;
         battery = true;
-        bi.updateTimeBasesLocked(battery, screen, clocks.realtime*1000, 0);
+        bi.updateTimeBasesLocked(battery, screen[0], clocks.realtime * 1000, 0);
         bi.setOnBatteryInternal(battery);
         clocks.realtime += 19;
-        bi.updateDisplayMeasuredEnergyStatsLocked(510_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{510_000}, screen, clocks.realtime);
         checkMeasuredCharge("B", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case C: uid1 ON,  uid2 off, battery on,  screen off
         clocks.realtime += 18;
         setFgState(uid1, true, bi);
         clocks.realtime += 18;
-        bi.updateDisplayMeasuredEnergyStatsLocked(520_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{520_000}, screen, clocks.realtime);
         checkMeasuredCharge("C", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case D: uid1 on,  uid2 off, battery on,  screen ON
         clocks.realtime += 17;
-        screen = Display.STATE_ON;
-        bi.updateDisplayMeasuredEnergyStatsLocked(521_000, screen, clocks.realtime);
+        screen[0] = Display.STATE_ON;
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{521_000}, screen, clocks.realtime);
         blame1 += 0; // Screen had been off during the measurement period
         checkMeasuredCharge("D.1", uid1, blame1, uid2, blame2, globalDoze, bi);
         clocks.realtime += 101;
-        bi.updateDisplayMeasuredEnergyStatsLocked(530_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{530_000}, screen, clocks.realtime);
         blame1 += 530_000;
         checkMeasuredCharge("D.2", uid1, blame1, uid2, blame2, globalDoze, bi);
 
@@ -641,33 +1044,33 @@ public class BatteryStatsNoteTest extends TestCase {
         clocks.realtime += 20;
         setFgState(uid2, true, bi);
         clocks.realtime += 40;
-        bi.updateDisplayMeasuredEnergyStatsLocked(540_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{540_000}, screen, clocks.realtime);
         // In the past 60ms, sum of fg is 20+40+40=100ms. uid1 is blamed for 60/100; uid2 for 40/100
         blame1 += 540_000 * (20 + 40) / (20 + 40 + 40);
-        blame2 += 540_000 * ( 0 + 40) / (20 + 40 + 40);
+        blame2 += 540_000 * (0 + 40) / (20 + 40 + 40);
         checkMeasuredCharge("E", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case F: uid1 on,  uid2 OFF, battery on,  screen on
         clocks.realtime += 40;
         setFgState(uid2, false, bi);
         clocks.realtime += 120;
-        bi.updateDisplayMeasuredEnergyStatsLocked(550_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{550_000}, screen, clocks.realtime);
         // In the past 160ms, sum f fg is 200ms. uid1 is blamed for 40+120 of it; uid2 for 40 of it.
         blame1 += 550_000 * (40 + 120) / (40 + 40 + 120);
-        blame2 += 550_000 * (40 + 0  ) / (40 + 40 + 120);
+        blame2 += 550_000 * (40 + 0) / (40 + 40 + 120);
         checkMeasuredCharge("F", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case G: uid1 on,  uid2 off,  battery on, screen DOZE
         clocks.realtime += 5;
-        screen = Display.STATE_DOZE;
-        bi.updateDisplayMeasuredEnergyStatsLocked(570_000, screen, clocks.realtime);
+        screen[0] = Display.STATE_DOZE;
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{570_000}, screen, clocks.realtime);
         blame1 += 570_000; // All of this pre-doze time is blamed on uid1.
         checkMeasuredCharge("G", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case H: uid1 on,  uid2 off,  battery on, screen ON
         clocks.realtime += 6;
-        screen = Display.STATE_ON;
-        bi.updateDisplayMeasuredEnergyStatsLocked(580_000, screen, clocks.realtime);
+        screen[0] = Display.STATE_ON;
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{580_000}, screen, clocks.realtime);
         blame1 += 0; // The screen had been doze during the energy period
         globalDoze += 580_000;
         checkMeasuredCharge("H", uid1, blame1, uid2, blame2, globalDoze, bi);
@@ -822,4 +1225,19 @@ public class BatteryStatsNoteTest extends TestCase {
 
         assertEquals("Wrong uid2 blame in bucket 1 for Case " + caseName, blame2B, actualUid2[1]);
     }
+
+    private void checkScreenBrightnesses(long[] overallExpected, long[][] perDisplayExpected,
+            BatteryStatsImpl bi, long currentTimeMs) {
+        final int numDisplay = bi.getDisplayCount();
+        for (int bin = 0; bin < NUM_SCREEN_BRIGHTNESS_BINS; bin++) {
+            for (int display = 0; display < numDisplay; display++) {
+                assertEquals("Failure for display " + display + " screen brightness bin " + bin,
+                        perDisplayExpected[display][bin] * 1000,
+                        bi.getDisplayScreenBrightnessTime(display, bin, currentTimeMs * 1000));
+            }
+            assertEquals("Failure for overall screen brightness bin " + bin,
+                    overallExpected[bin] * 1000,
+                    bi.getScreenBrightnessTime(bin, currentTimeMs * 1000, STATS_SINCE_CHARGED));
+        }
+    }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsRule.java b/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsRule.java
index 083090c54619..ac87806b1639 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsRule.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsRule.java
@@ -110,6 +110,20 @@ public class BatteryUsageStatsRule implements TestRule {
         return this;
     }
 
+    public BatteryUsageStatsRule setAveragePowerForOrdinal(String group, int ordinal,
+            double value) {
+        when(mPowerProfile.getAveragePowerForOrdinal(group, ordinal)).thenReturn(value);
+        when(mPowerProfile.getAveragePowerForOrdinal(eq(group), eq(ordinal),
+                anyDouble())).thenReturn(value);
+        return this;
+    }
+
+    public BatteryUsageStatsRule setNumDisplays(int value) {
+        when(mPowerProfile.getNumDisplays()).thenReturn(value);
+        mBatteryStats.setDisplayCountLocked(value);
+        return this;
+    }
+
     /** Call only after setting the power profile information. */
     public BatteryUsageStatsRule initMeasuredEnergyStatsLocked() {
         return initMeasuredEnergyStatsLocked(new String[0]);
diff --git a/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java b/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
index cee1a0352a7e..cfecf15b55ef 100644
--- a/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
+++ b/core/tests/coretests/src/com/android/internal/os/MockBatteryStatsImpl.java
@@ -39,13 +39,14 @@ public class MockBatteryStatsImpl extends BatteryStatsImpl {
     public BatteryStatsImpl.Clocks clocks;
     public boolean mForceOnBattery;
     private NetworkStats mNetworkStats;
+    private DummyExternalStatsSync mExternalStatsSync = new DummyExternalStatsSync();
 
     MockBatteryStatsImpl(Clocks clocks) {
         super(clocks);
         this.clocks = mClocks;
         initTimersAndCounters();
 
-        setExternalStatsSyncLocked(new DummyExternalStatsSync());
+        setExternalStatsSyncLocked(mExternalStatsSync);
         informThatAllExternalStatsAreFlushed();
 
         // A no-op handler.
@@ -182,7 +183,15 @@ public class MockBatteryStatsImpl extends BatteryStatsImpl {
         return mPendingUids;
     }
 
+    public int getAndClearExternalStatsSyncFlags() {
+        final int flags = mExternalStatsSync.flags;
+        mExternalStatsSync.flags = 0;
+        return flags;
+    }
+
     private class DummyExternalStatsSync implements ExternalStatsSync {
+        public int flags = 0;
+
         @Override
         public Future<?> scheduleSync(String reason, int flags) {
             return null;
@@ -211,8 +220,9 @@ public class MockBatteryStatsImpl extends BatteryStatsImpl {
         }
 
         @Override
-        public Future<?> scheduleSyncDueToScreenStateChange(
-                int flag, boolean onBattery, boolean onBatteryScreenOff, int screenState) {
+        public Future<?> scheduleSyncDueToScreenStateChange(int flag, boolean onBattery,
+                boolean onBatteryScreenOff, int screenState, int[] perDisplayScreenStates) {
+            flags |= flag;
             return null;
         }
 
diff --git a/core/tests/coretests/src/com/android/internal/os/PowerProfileTest.java b/core/tests/coretests/src/com/android/internal/os/PowerProfileTest.java
index 5862368f44d2..88ee405483db 100644
--- a/core/tests/coretests/src/com/android/internal/os/PowerProfileTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/PowerProfileTest.java
@@ -17,6 +17,10 @@
 
 package com.android.internal.os;
 
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_AMBIENT;
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_FULL;
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON;
+
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
 
@@ -53,7 +57,12 @@ public class PowerProfileTest extends TestCase {
         assertEquals(4, mProfile.getNumSpeedStepsInCpuCluster(1));
         assertEquals(60.0, mProfile.getAveragePowerForCpuCore(1, 3));
         assertEquals(3000.0, mProfile.getBatteryCapacity());
-        assertEquals(0.5, mProfile.getAveragePower(PowerProfile.POWER_AMBIENT_DISPLAY));
+        assertEquals(0.5,
+                mProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 0));
+        assertEquals(100.0,
+                mProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, 0));
+        assertEquals(800.0,
+                mProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 0));
         assertEquals(100.0, mProfile.getAveragePower(PowerProfile.POWER_AUDIO));
         assertEquals(150.0, mProfile.getAveragePower(PowerProfile.POWER_VIDEO));
     }
diff --git a/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
index c695fc9eb87d..eee5d57c7bc6 100644
--- a/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
@@ -16,6 +16,9 @@
 
 package com.android.internal.os;
 
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_FULL;
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import android.app.ActivityManager;
@@ -39,24 +42,27 @@ public class ScreenPowerCalculatorTest {
     private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 43;
     private static final long MINUTE_IN_MS = 60 * 1000;
     private static final long MINUTE_IN_US = 60 * 1000 * 1000;
+    private static final long HOUR_IN_MS = 60 * MINUTE_IN_MS;
 
     @Rule
     public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
-            .setAveragePower(PowerProfile.POWER_SCREEN_ON, 36.0)
-            .setAveragePower(PowerProfile.POWER_SCREEN_FULL, 48.0);
+            .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, 0, 36.0)
+            .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 0, 48.0)
+            .setNumDisplays(1);
 
     @Test
     public void testMeasuredEnergyBasedModel() {
         mStatsRule.initMeasuredEnergyStatsLocked();
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
 
-        batteryStats.noteScreenStateLocked(Display.STATE_ON, 0, 0, 0);
-        batteryStats.updateDisplayMeasuredEnergyStatsLocked(0, Display.STATE_ON, 0);
+        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, 0, 0, 0);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{0},
+                new int[]{Display.STATE_ON}, 0);
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true,
                 0, 0);
 
-        batteryStats.updateDisplayMeasuredEnergyStatsLocked(200_000_000, Display.STATE_ON,
-                15 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{200_000_000},
+                new int[]{Display.STATE_ON}, 15 * MINUTE_IN_MS);
 
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_CACHED_EMPTY, false,
                 20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
@@ -64,16 +70,16 @@ public class ScreenPowerCalculatorTest {
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
                 20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
 
-        batteryStats.updateDisplayMeasuredEnergyStatsLocked(300_000_000, Display.STATE_ON,
-                60 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{300_000_000},
+                new int[]{Display.STATE_ON}, 60 * MINUTE_IN_MS);
 
-        batteryStats.noteScreenStateLocked(Display.STATE_OFF,
+        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
 
-        batteryStats.updateDisplayMeasuredEnergyStatsLocked(100_000_000, Display.STATE_DOZE,
-                120 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{100_000_000},
+                new int[]{Display.STATE_DOZE}, 120 * MINUTE_IN_MS);
 
         mStatsRule.setTime(120 * MINUTE_IN_US, 120 * MINUTE_IN_US);
 
@@ -126,24 +132,122 @@ public class ScreenPowerCalculatorTest {
                 .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
     }
 
+
+    @Test
+    public void testMeasuredEnergyBasedModel_multiDisplay() {
+        mStatsRule.initMeasuredEnergyStatsLocked()
+                .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, 1, 60.0)
+                .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 1, 100.0)
+                .setNumDisplays(2);
+
+        BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+
+        final int[] screenStates = new int[]{Display.STATE_ON, Display.STATE_OFF};
+
+        batteryStats.noteScreenStateLocked(0, screenStates[0], 0, 0, 0);
+        batteryStats.noteScreenStateLocked(1, screenStates[1], 0, 0, 0);
+        batteryStats.noteScreenBrightnessLocked(0, 255, 0, 0);
+        setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true, 0, 0);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{300, 400}, screenStates, 0);
+
+        batteryStats.noteScreenBrightnessLocked(0, 100, 5 * MINUTE_IN_MS, 5 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(0, 200, 10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
+
+        setProcState(APP_UID1, ActivityManager.PROCESS_STATE_CACHED_EMPTY, false,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+        setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+
+        screenStates[0] = Display.STATE_OFF;
+        screenStates[1] = Display.STATE_ON;
+        batteryStats.noteScreenStateLocked(0, screenStates[0],
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, screenStates[1], 80 * MINUTE_IN_MS,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{600_000_000, 500},
+                screenStates, 80 * MINUTE_IN_MS);
+
+        batteryStats.noteScreenBrightnessLocked(1, 25, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(1, 250, 86 * MINUTE_IN_MS, 86 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(1, 75, 98 * MINUTE_IN_MS, 98 * MINUTE_IN_MS);
+
+        screenStates[1] = Display.STATE_OFF;
+        batteryStats.noteScreenStateLocked(1, screenStates[1], 110 * MINUTE_IN_MS,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{700, 800_000_000},
+                screenStates, 110 * MINUTE_IN_MS);
+
+        setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+
+        mStatsRule.setTime(120 * MINUTE_IN_US, 120 * MINUTE_IN_US);
+
+        ScreenPowerCalculator calculator =
+                new ScreenPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(calculator);
+
+        BatteryConsumer deviceConsumer = mStatsRule.getDeviceBatteryConsumer();
+        assertThat(deviceConsumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(110 * MINUTE_IN_MS);
+        // (600000000 + 800000000) uAs * (1 mA / 1000 uA) * (1 h / 3600 s)  = 166.66666 mAh
+        assertThat(deviceConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(388.88888);
+        assertThat(deviceConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+
+        UidBatteryConsumer uid1 = mStatsRule.getUidBatteryConsumer(APP_UID1);
+        assertThat(uid1.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(20 * MINUTE_IN_MS);
+
+        // Uid1 ran for 20 out of 80 min during the first Display update.
+        // It also ran for 5 out of 45 min during the second Display update:
+        // Uid1 charge = 20 / 80 * 600000000 mAs = 41.66666 mAh
+        assertThat(uid1.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(41.66666);
+        assertThat(uid1.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+
+        UidBatteryConsumer uid2 = mStatsRule.getUidBatteryConsumer(APP_UID2);
+        assertThat(uid2.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(90 * MINUTE_IN_MS);
+
+        // Uid2 ran for 60 out of 80 min during the first Display update.
+        // It also ran for all of the second Display update:
+        // Uid1 charge = 60 / 80 * 600000000 + 800000000 mAs = 347.22222 mAh
+        assertThat(uid2.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(347.22222);
+        assertThat(uid2.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+
+        BatteryConsumer appsConsumer = mStatsRule.getAppsBatteryConsumer();
+        assertThat(appsConsumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(110 * MINUTE_IN_MS);
+        assertThat(appsConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(388.88888);
+        assertThat(appsConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+
+    }
+
     @Test
     public void testPowerProfileBasedModel() {
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
 
-        batteryStats.noteScreenStateLocked(Display.STATE_ON, 0, 0, 0);
-        batteryStats.noteScreenBrightnessLocked(255, 0, 0);
+        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, 0, 0, 0);
+        batteryStats.noteScreenBrightnessLocked(0, 255, 0, 0);
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true,
                 0, 0);
 
-        batteryStats.noteScreenBrightnessLocked(100, 5 * MINUTE_IN_MS, 5 * MINUTE_IN_MS);
-        batteryStats.noteScreenBrightnessLocked(200, 10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(0, 100, 5 * MINUTE_IN_MS, 5 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(0, 200, 10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
 
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_CACHED_EMPTY, false,
                 20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
                 20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
 
-        batteryStats.noteScreenStateLocked(Display.STATE_OFF,
+        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
@@ -194,6 +298,95 @@ public class ScreenPowerCalculatorTest {
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
     }
 
+
+    @Test
+    public void testPowerProfileBasedModel_multiDisplay() {
+        mStatsRule.setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, 1, 60.0)
+                .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 1, 100.0)
+                .setNumDisplays(2);
+
+        BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+
+        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, 0, 0, 0);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_OFF, 0, 0, 0);
+        batteryStats.noteScreenBrightnessLocked(0, 255, 0, 0);
+        setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true,
+                0, 0);
+
+        batteryStats.noteScreenBrightnessLocked(0, 100, 5 * MINUTE_IN_MS, 5 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(0, 200, 10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
+
+        setProcState(APP_UID1, ActivityManager.PROCESS_STATE_CACHED_EMPTY, false,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+        setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+
+        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_ON, 80 * MINUTE_IN_MS,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(1, 20, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+
+        batteryStats.noteScreenBrightnessLocked(1, 250, 86 * MINUTE_IN_MS, 86 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(1, 75, 98 * MINUTE_IN_MS, 98 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_OFF, 110 * MINUTE_IN_MS,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+
+        setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+
+        mStatsRule.setTime(120 * MINUTE_IN_US, 120 * MINUTE_IN_US);
+        ScreenPowerCalculator calculator =
+                new ScreenPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(BatteryUsageStatsRule.POWER_PROFILE_MODEL_ONLY, calculator);
+
+        BatteryConsumer deviceConsumer = mStatsRule.getDeviceBatteryConsumer();
+        assertThat(deviceConsumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(110 * MINUTE_IN_MS);
+        // First display consumed 92 mAh.
+        // Second display ran for 0.5 hours at a base drain rate of 60 mA.
+        // 6 minutes (0.1 hours) spent in the first brightness level which drains an extra 10 mA.
+        // 12 minutes (0.2 hours) spent in the fifth brightness level which drains an extra 90 mA.
+        // 12 minutes (0.2 hours) spent in the second brightness level which drains an extra 30 mA.
+        // 92 + 60 * 0.5 + 10 * 0.1 + 90 * 0.2 + 30 * 0.2 = 147
+        assertThat(deviceConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(147);
+        assertThat(deviceConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+
+        UidBatteryConsumer uid1 = mStatsRule.getUidBatteryConsumer(APP_UID1);
+        assertThat(uid1.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(20 * MINUTE_IN_MS);
+
+        // Uid1 took 20 out of the total of 110 min of foreground activity
+        // Uid1 charge = 20 / 110 * 147.0 = 23.0 mAh
+        assertThat(uid1.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(26.72727);
+        assertThat(uid1.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+
+        UidBatteryConsumer uid2 = mStatsRule.getUidBatteryConsumer(APP_UID2);
+        assertThat(uid2.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(90 * MINUTE_IN_MS);
+
+        // Uid2 took 90 out of the total of 110 min of foreground activity
+        // Uid2 charge = 90 / 110 * 92.0 = 69.0 mAh
+        assertThat(uid2.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(120.272727);
+        assertThat(uid2.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+
+        BatteryConsumer appsConsumer = mStatsRule.getAppsBatteryConsumer();
+        assertThat(appsConsumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(110 * MINUTE_IN_MS);
+        assertThat(appsConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(147);
+        assertThat(appsConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+
+    }
+
     private void setProcState(int uid, int procState, boolean resumed, long realtimeMs,
             long uptimeMs) {
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
diff --git a/core/tests/coretests/src/com/android/internal/statusbar/RegisterStatusBarResultTest.java b/core/tests/coretests/src/com/android/internal/statusbar/RegisterStatusBarResultTest.java
index 272f2287dd6e..0f05be06bff6 100644
--- a/core/tests/coretests/src/com/android/internal/statusbar/RegisterStatusBarResultTest.java
+++ b/core/tests/coretests/src/com/android/internal/statusbar/RegisterStatusBarResultTest.java
@@ -24,6 +24,7 @@ import android.os.Binder;
 import android.os.Parcel;
 import android.os.UserHandle;
 import android.util.ArrayMap;
+import android.view.InsetsVisibilities;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -59,7 +60,8 @@ public class RegisterStatusBarResultTest {
                 new Binder() /* imeToken */,
                 true /* navbarColorManagedByIme */,
                 BEHAVIOR_SHOW_TRANSIENT_BARS_BY_SWIPE,
-                true /* appFullscreen */,
+                new InsetsVisibilities() /* requestedVisibilities */,
+                "test" /* packageName */,
                 new int[0] /* transientBarTypes */);
 
         final RegisterStatusBarResult copy = clone(original);
@@ -79,7 +81,8 @@ public class RegisterStatusBarResultTest {
         assertThat(copy.mImeToken).isSameInstanceAs(original.mImeToken);
         assertThat(copy.mNavbarColorManagedByIme).isEqualTo(original.mNavbarColorManagedByIme);
         assertThat(copy.mBehavior).isEqualTo(original.mBehavior);
-        assertThat(copy.mAppFullscreen).isEqualTo(original.mAppFullscreen);
+        assertThat(copy.mRequestedVisibilities).isEqualTo(original.mRequestedVisibilities);
+        assertThat(copy.mPackageName).isEqualTo(original.mPackageName);
         assertThat(copy.mTransientBarTypes).isEqualTo(original.mTransientBarTypes);
     }
 
diff --git a/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateManagerGlobalTest.java b/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateManagerGlobalTest.java
index 055fc7171fa0..db63e6e0b187 100644
--- a/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateManagerGlobalTest.java
+++ b/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateManagerGlobalTest.java
@@ -160,6 +160,34 @@ public final class DeviceStateManagerGlobalTest {
         verify(callback).onStateChanged(eq(mService.getBaseState()));
     }
 
+    @Test
+    public void verifyDeviceStateRequestCallbacksCalled() {
+        DeviceStateRequest.Callback callback = mock(TestDeviceStateRequestCallback.class);
+
+        DeviceStateRequest request = DeviceStateRequest.newBuilder(OTHER_DEVICE_STATE).build();
+        mDeviceStateManagerGlobal.requestState(request,
+                ConcurrentUtils.DIRECT_EXECUTOR /* executor */,
+                callback /* callback */);
+
+        verify(callback).onRequestActivated(eq(request));
+        Mockito.reset(callback);
+
+        mDeviceStateManagerGlobal.cancelRequest(request);
+
+        verify(callback).onRequestCanceled(eq(request));
+    }
+
+    public static class TestDeviceStateRequestCallback implements DeviceStateRequest.Callback {
+        @Override
+        public void onRequestActivated(DeviceStateRequest request) { }
+
+        @Override
+        public void onRequestCanceled(DeviceStateRequest request) { }
+
+        @Override
+        public void onRequestSuspended(DeviceStateRequest request) { }
+    }
+
     private static final class TestDeviceStateManagerService extends IDeviceStateManager.Stub {
         public static final class Request {
             public final IBinder token;
diff --git a/core/tests/mockingcoretests/src/android/window/ConfigurationHelperTest.java b/core/tests/mockingcoretests/src/android/window/ConfigurationHelperTest.java
new file mode 100644
index 000000000000..996d7b435e5a
--- /dev/null
+++ b/core/tests/mockingcoretests/src/android/window/ConfigurationHelperTest.java
@@ -0,0 +1,164 @@
+/*
+ * Copyright (C) 2021 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 android.window;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+
+import android.app.ResourcesManager;
+import android.content.res.Configuration;
+import android.graphics.Rect;
+import android.os.Binder;
+import android.platform.test.annotations.Presubmit;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoSession;
+import org.mockito.quality.Strictness;
+
+/**
+ * Tests for {@link ConfigurationHelper}
+ *
+ * <p>Build/Install/Run:
+ *  atest FrameworksMockingCoreTests:ConfigurationHelperTest
+ *
+ * <p>This test class is a part of Window Manager Service tests and specified in
+ * {@link com.android.server.wm.test.filters.FrameworksTestsFilter}.
+ */
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+@Presubmit
+public class ConfigurationHelperTest {
+    MockitoSession mMockitoSession;
+    ResourcesManager mResourcesManager;
+
+    @Before
+    public void setUp() {
+        mMockitoSession = mockitoSession()
+                .strictness(Strictness.LENIENT)
+                .spyStatic(ResourcesManager.class)
+                .startMocking();
+        doReturn(mock(ResourcesManager.class)).when(ResourcesManager::getInstance);
+        mResourcesManager = ResourcesManager.getInstance();
+    }
+
+    @After
+    public void tearDown() {
+        mMockitoSession.finishMocking();
+    }
+
+    @Test
+    public void testShouldUpdateResources_NullConfig_ReturnsTrue() {
+        assertThat(ConfigurationHelper.shouldUpdateResources(new Binder(), null /* config */,
+                new Configuration(), new Configuration(), false /* displayChanged */,
+                null /* configChanged */)).isTrue();
+    }
+
+    @Test
+    public void testShouldUpdateResources_DisplayChanged_ReturnsTrue() {
+        assertThat(ConfigurationHelper.shouldUpdateResources(new Binder(), new Configuration(),
+                new Configuration(), new Configuration(), true /* displayChanged */,
+                null /* configChanged */)).isTrue();
+    }
+
+    @Test
+    public void testShouldUpdateResources_DifferentResources_ReturnsTrue() {
+        doReturn(false).when(mResourcesManager).isSameResourcesOverrideConfig(any(), any());
+
+        assertThat(ConfigurationHelper.shouldUpdateResources(new Binder(), new Configuration(),
+                new Configuration(), new Configuration(), false /* displayChanged */,
+                null /* configChanged */)).isTrue();
+    }
+
+    @Test
+    public void testShouldUpdateResources_DifferentBounds_ReturnsTrue() {
+        doReturn(true).when(mResourcesManager).isSameResourcesOverrideConfig(any(), any());
+
+        final Configuration config = new Configuration();
+        config.windowConfiguration.setBounds(new Rect(0, 0, 10, 10));
+        config.windowConfiguration.setMaxBounds(new Rect(0, 0, 20, 20));
+
+        final Configuration newConfig = new Configuration();
+        newConfig.windowConfiguration.setBounds(new Rect(0, 0, 20, 20));
+        newConfig.windowConfiguration.setMaxBounds(new Rect(0, 0, 20, 20));
+
+        assertThat(ConfigurationHelper.shouldUpdateResources(new Binder(), config, newConfig,
+                new Configuration(), false /* displayChanged */, null /* configChanged */))
+                .isTrue();
+    }
+
+    @Test
+    public void testShouldUpdateResources_SameConfig_ReturnsFalse() {
+        doReturn(true).when(mResourcesManager).isSameResourcesOverrideConfig(any(), any());
+
+        final Configuration config = new Configuration();
+        final Configuration newConfig = new Configuration();
+
+        assertThat(ConfigurationHelper.shouldUpdateResources(new Binder(), config, newConfig,
+                new Configuration(), false /* displayChanged */, null /* configChanged */))
+                .isFalse();
+    }
+
+    @Test
+    public void testShouldUpdateResources_DifferentConfig_ReturnsTrue() {
+        doReturn(true).when(mResourcesManager).isSameResourcesOverrideConfig(any(), any());
+
+        final Configuration config = new Configuration();
+        final Configuration newConfig = new Configuration();
+        newConfig.setToDefaults();
+
+        assertThat(ConfigurationHelper.shouldUpdateResources(new Binder(), config, newConfig,
+                new Configuration(), false /* displayChanged */, null /* configChanged */))
+                .isTrue();
+    }
+
+    @Test
+    public void testShouldUpdateResources_DifferentNonPublicConfig_ReturnsTrue() {
+        doReturn(true).when(mResourcesManager).isSameResourcesOverrideConfig(any(), any());
+
+        final Configuration config = new Configuration();
+        final Configuration newConfig = new Configuration();
+        newConfig.windowConfiguration.setAppBounds(new Rect(0, 0, 10, 10));
+
+        assertThat(ConfigurationHelper.shouldUpdateResources(new Binder(), config, newConfig,
+                new Configuration(), false /* displayChanged */, null /* configChanged */))
+                .isTrue();
+    }
+
+    @Test
+    public void testShouldUpdateResources_OverrideConfigChanged_ReturnsFalse() {
+        doReturn(true).when(mResourcesManager).isSameResourcesOverrideConfig(any(), any());
+
+        final Configuration config = new Configuration();
+        final Configuration newConfig = new Configuration();
+        final boolean configChanged = true;
+
+        assertThat(ConfigurationHelper.shouldUpdateResources(new Binder(), config, newConfig,
+                new Configuration(), false /* displayChanged */, configChanged))
+                .isEqualTo(configChanged);
+    }
+}
diff --git a/core/tests/mockingcoretests/src/android/window/SizeConfigurationBucketsTest.java b/core/tests/mockingcoretests/src/android/window/SizeConfigurationBucketsTest.java
new file mode 100644
index 000000000000..fa4aa803c75e
--- /dev/null
+++ b/core/tests/mockingcoretests/src/android/window/SizeConfigurationBucketsTest.java
@@ -0,0 +1,379 @@
+/*
+ * Copyright (C) 2021 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 android.window;
+
+import static android.content.pm.ActivityInfo.CONFIG_LOCALE;
+import static android.content.pm.ActivityInfo.CONFIG_SCREEN_LAYOUT;
+import static android.content.pm.ActivityInfo.CONFIG_SCREEN_SIZE;
+import static android.content.pm.ActivityInfo.CONFIG_SMALLEST_SCREEN_SIZE;
+import static android.content.res.Configuration.SCREENLAYOUT_COMPAT_NEEDED;
+import static android.content.res.Configuration.SCREENLAYOUT_LAYOUTDIR_LTR;
+import static android.content.res.Configuration.SCREENLAYOUT_LAYOUTDIR_RTL;
+import static android.content.res.Configuration.SCREENLAYOUT_LAYOUTDIR_UNDEFINED;
+import static android.content.res.Configuration.SCREENLAYOUT_LONG_NO;
+import static android.content.res.Configuration.SCREENLAYOUT_LONG_YES;
+import static android.content.res.Configuration.SCREENLAYOUT_ROUND_NO;
+import static android.content.res.Configuration.SCREENLAYOUT_ROUND_UNDEFINED;
+import static android.content.res.Configuration.SCREENLAYOUT_ROUND_YES;
+import static android.content.res.Configuration.SCREENLAYOUT_SIZE_LARGE;
+import static android.content.res.Configuration.SCREENLAYOUT_SIZE_NORMAL;
+import static android.content.res.Configuration.SCREENLAYOUT_SIZE_SMALL;
+import static android.content.res.Configuration.SCREENLAYOUT_SIZE_XLARGE;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import android.content.res.Configuration;
+import android.platform.test.annotations.Presubmit;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Arrays;
+
+/**
+ * Tests for {@link SizeConfigurationBuckets}
+ *
+ * Build/Install/Run:
+ *  atest FrameworksMockingCoreTests:SizeConfigurationBucketsTest
+ */
+
+@SmallTest
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class SizeConfigurationBucketsTest {
+
+    /**
+     * Tests that a change in any of the non-size-related screen layout fields results in
+     * {@link SizeConfigurationBuckets#areNonSizeLayoutFieldsUnchanged} returning false.
+     */
+    @Test
+    public void testNonSizeRelatedScreenLayoutFields() {
+        // Test layout direction
+        assertEquals(true, SizeConfigurationBuckets
+                .areNonSizeLayoutFieldsUnchanged(0, SCREENLAYOUT_LAYOUTDIR_UNDEFINED));
+        assertEquals(false, SizeConfigurationBuckets
+                .areNonSizeLayoutFieldsUnchanged(0, SCREENLAYOUT_LAYOUTDIR_LTR));
+        assertEquals(false, SizeConfigurationBuckets
+                .areNonSizeLayoutFieldsUnchanged(0, SCREENLAYOUT_LAYOUTDIR_RTL));
+
+        // Test layout roundness
+        assertEquals(true, SizeConfigurationBuckets
+                .areNonSizeLayoutFieldsUnchanged(0, SCREENLAYOUT_ROUND_UNDEFINED));
+        assertEquals(false, SizeConfigurationBuckets
+                .areNonSizeLayoutFieldsUnchanged(0, SCREENLAYOUT_ROUND_NO));
+        assertEquals(false, SizeConfigurationBuckets
+                .areNonSizeLayoutFieldsUnchanged(0, SCREENLAYOUT_ROUND_YES));
+
+        // Test layout compat needed
+        assertEquals(false, SizeConfigurationBuckets
+                .areNonSizeLayoutFieldsUnchanged(0, SCREENLAYOUT_COMPAT_NEEDED));
+    }
+
+    /**
+     * Tests that null size configuration buckets unflips the correct configuration flags.
+     */
+    @Test
+    public void testNullSizeConfigurationBuckets() {
+        // Check that all 3 size configurations are filtered out of the diff if the buckets are null
+        // and non-size attributes of screen layout are unchanged. Add a non-size related config
+        // change (i.e. CONFIG_LOCALE) to test that the diff is not set to zero.
+        final int diff = CONFIG_SCREEN_SIZE | CONFIG_SMALLEST_SCREEN_SIZE | CONFIG_SCREEN_LAYOUT
+                | CONFIG_LOCALE;
+        final int filteredDiffNonSizeLayoutUnchanged = SizeConfigurationBuckets.filterDiff(diff,
+                Configuration.EMPTY, Configuration.EMPTY, null);
+        assertEquals(CONFIG_LOCALE, filteredDiffNonSizeLayoutUnchanged);
+
+        // Check that only screen size and smallest screen size are filtered out of the diff if the
+        // buckets are null and non-size attributes of screen layout are changed.
+        final Configuration newConfig = new Configuration();
+        newConfig.screenLayout |= SCREENLAYOUT_ROUND_YES;
+        final int filteredDiffNonSizeLayoutChanged = SizeConfigurationBuckets.filterDiff(diff,
+                Configuration.EMPTY, newConfig, null);
+        assertEquals(CONFIG_SCREEN_LAYOUT | CONFIG_LOCALE, filteredDiffNonSizeLayoutChanged);
+    }
+
+    /**
+     * Tests that {@link SizeConfigurationBuckets.crossesSizeThreshold()} correctly checks whether
+     * the bucket thresholds have or have not been crossed. This test includes boundary checks
+     * to ensure that arithmetic is inclusive and exclusive in the right places.
+     */
+    @Test
+    public void testCrossesSizeThreshold() {
+        final int[] thresholds = new int[] { 360, 600 };
+        final int nThresholds = thresholds.length;
+        for (int i = -1; i < nThresholds; i++) {
+            final int minValueInBucket = i < 0 ? 0 : thresholds[i];
+            final int maxValueInBucket = i < nThresholds - 1
+                    ? thresholds[i + 1] - 1 : Integer.MAX_VALUE;
+            final int bucketRange = maxValueInBucket - minValueInBucket;
+            // Set old value to 1/4 in between the two thresholds.
+            final int oldValue = (int) (minValueInBucket + bucketRange * 0.25);
+            // Test 3 values of new value spread across bucket range: minValueInBucket, bucket
+            // midpoint, and max value in bucket. In all 3 cases, the bucket has not changed so
+            // {@link SizeConfigurationBuckets#crossedSizeThreshold()} should return false.
+            checkCrossesSizeThreshold(thresholds, oldValue, minValueInBucket, false);
+            checkCrossesSizeThreshold(thresholds, oldValue,
+                    (int) (minValueInBucket + bucketRange * 0.5), false);
+            checkCrossesSizeThreshold(thresholds, oldValue, maxValueInBucket, false);
+            // Test 4 values of size spread outside of bucket range: more than 1 less than min
+            // value, 1 less than min value, 1 more than max value, and more than 1 more than max
+            // value. In all 4 cases, the bucket has changed so
+            // {@link SizeConfigurationBuckets#crossedSizeThreshold()} should return true.
+            // Only test less than min value if min value > 0.
+            if (minValueInBucket > 0) {
+                checkCrossesSizeThreshold(thresholds, oldValue, minValueInBucket - 20, true);
+                checkCrossesSizeThreshold(thresholds, oldValue, minValueInBucket - 1, true);
+            }
+            // Only test greater than max value if not in highest bucket.
+            if (i < nThresholds - 1) {
+                checkCrossesSizeThreshold(thresholds, oldValue, maxValueInBucket + 1, true);
+                checkCrossesSizeThreshold(thresholds, oldValue, maxValueInBucket + 20, true);
+            }
+        }
+    }
+
+    /**
+     * Tests that if screen layout size changed but did not cross a threshold, the filtered diff
+     * does not include screen layout.
+     */
+    @Test
+    public void testScreenLayoutFilteredIfSizeDidNotCrossThreshold() {
+        // Set only small and large sizes
+        final Configuration[] sizeConfigs = new Configuration[2];
+        sizeConfigs[0] = new Configuration();
+        sizeConfigs[0].screenLayout |= SCREENLAYOUT_SIZE_SMALL;
+        sizeConfigs[1] = new Configuration();
+        sizeConfigs[1].screenLayout |= SCREENLAYOUT_SIZE_LARGE;
+        final SizeConfigurationBuckets sizeBuckets = new SizeConfigurationBuckets(sizeConfigs);
+
+        // Change screen layout size from small to normal and check that screen layout flag is
+        // not part of the diff because a threshold was not crossed.
+        final int diff = CONFIG_SCREEN_LAYOUT;
+        final Configuration oldConfig = new Configuration();
+        oldConfig.screenLayout |= SCREENLAYOUT_SIZE_SMALL;
+        final Configuration newConfig = new Configuration();
+        newConfig.screenLayout |= SCREENLAYOUT_SIZE_NORMAL;
+        final int filteredDiff = SizeConfigurationBuckets.filterDiff(diff, oldConfig, newConfig,
+                sizeBuckets);
+        assertEquals(0, filteredDiff);
+
+        // If a non-size attribute of screen layout changed, then screen layout should not be
+        // filtered from the diff.
+        newConfig.screenLayout |= SCREENLAYOUT_ROUND_YES;
+        final int filteredDiffNonSizeLayoutChanged = SizeConfigurationBuckets.filterDiff(diff,
+                oldConfig, newConfig, sizeBuckets);
+        assertEquals(CONFIG_SCREEN_LAYOUT, filteredDiffNonSizeLayoutChanged);
+    }
+
+    /**
+     * Tests that if screen layout size changed and did cross a threshold, the filtered diff
+     * includes screen layout.
+     */
+    @Test
+    public void testScreenLayoutNotFilteredIfSizeCrossedThreshold() {
+        // Set only small and normal sizes
+        final Configuration[] sizeConfigs = new Configuration[2];
+        sizeConfigs[0] = new Configuration();
+        sizeConfigs[0].screenLayout |= SCREENLAYOUT_SIZE_SMALL;
+        sizeConfigs[1] = new Configuration();
+        sizeConfigs[1].screenLayout |= SCREENLAYOUT_SIZE_NORMAL;
+        final SizeConfigurationBuckets sizeBuckets = new SizeConfigurationBuckets(sizeConfigs);
+
+        // Change screen layout size from small to normal and check that screen layout flag is
+        // still part of the diff because a threshold was crossed.
+        final int diff = CONFIG_SCREEN_LAYOUT;
+        final Configuration oldConfig = new Configuration();
+        oldConfig.screenLayout |= SCREENLAYOUT_SIZE_SMALL;
+        final Configuration newConfig = new Configuration();
+        newConfig.screenLayout |= SCREENLAYOUT_SIZE_NORMAL;
+        final int filteredDiff = SizeConfigurationBuckets.filterDiff(diff, oldConfig, newConfig,
+                sizeBuckets);
+        assertEquals(CONFIG_SCREEN_LAYOUT, filteredDiff);
+    }
+
+    /**
+     * Tests that anytime screen layout size is decreased, the filtered diff still includes screen
+     * layout.
+     */
+    @Test
+    public void testScreenLayoutNotFilteredIfSizeDecreased() {
+        // The size thresholds can be anything, but can't be null
+        final int[] horizontalThresholds = new int[] { 360, 600 };
+        final SizeConfigurationBuckets sizeBuckets = new SizeConfigurationBuckets(
+                horizontalThresholds, null /* vertical */, null /* smallest */,
+                null /* screenLayoutSize */, false /* screenLayoutLongSet */);
+        final int[] sizeValuesInOrder = new int[] {
+                SCREENLAYOUT_SIZE_SMALL, SCREENLAYOUT_SIZE_NORMAL, SCREENLAYOUT_SIZE_LARGE,
+                SCREENLAYOUT_SIZE_XLARGE
+        };
+        final int nSizes = sizeValuesInOrder.length;
+        for (int larger = nSizes - 1; larger > 0; larger--) {
+            for (int smaller = larger - 1; smaller >= 0; smaller--) {
+                final Configuration oldConfig = new Configuration();
+                oldConfig.screenLayout |= sizeValuesInOrder[larger];
+                final Configuration newConfig = new Configuration();
+                newConfig.screenLayout |= sizeValuesInOrder[smaller];
+                assertTrue(String.format("oldSize=%d, newSize=%d", oldConfig.screenLayout,
+                        newConfig.screenLayout),
+                        sizeBuckets.crossesScreenLayoutSizeThreshold(oldConfig, newConfig));
+            }
+        }
+    }
+
+    /**
+     * Tests that if screen layout long changed but did not cross a threshold, the filtered diff
+     * does not include screen layout.
+     */
+    @Test
+    public void testScreenLayoutFilteredIfLongDidNotCrossThreshold() {
+        // Do not set any long threshold
+        final Configuration[] sizeConfigs = new Configuration[1];
+        sizeConfigs[0] = Configuration.EMPTY;
+        final SizeConfigurationBuckets sizeBuckets = new SizeConfigurationBuckets(sizeConfigs);
+
+        // Change screen layout long from not long to long and check that screen layout flag is
+        // not part of the diff because a threshold was not crossed.
+        final int diff = CONFIG_SCREEN_LAYOUT;
+        final Configuration oldConfig = new Configuration();
+        oldConfig.screenLayout |= SCREENLAYOUT_LONG_NO;
+        final Configuration newConfig = new Configuration();
+        newConfig.screenLayout |= SCREENLAYOUT_LONG_YES;
+        final int filteredDiff = SizeConfigurationBuckets.filterDiff(diff, oldConfig, newConfig,
+                sizeBuckets);
+        assertEquals(0, filteredDiff);
+
+        // If a non-size attribute of screen layout changed, then screen layout should not be
+        // filtered from the diff.
+        newConfig.screenLayout |= SCREENLAYOUT_ROUND_YES;
+        final int filteredDiffNonSizeLayoutChanged = SizeConfigurationBuckets.filterDiff(diff,
+                oldConfig, newConfig, sizeBuckets);
+        assertEquals(CONFIG_SCREEN_LAYOUT, filteredDiffNonSizeLayoutChanged);
+    }
+
+    /**
+     * Tests that if screen layout long changed and did cross a threshold, the filtered diff
+     * includes screen layout.
+     */
+    @Test
+    public void testScreenLayoutNotFilteredIfLongCrossedThreshold() {
+        // Set only small and normal sizes
+        final Configuration[] sizeConfigs = new Configuration[1];
+        sizeConfigs[0] = new Configuration();
+        sizeConfigs[0].screenLayout |= SCREENLAYOUT_LONG_NO;
+        final SizeConfigurationBuckets sizeBuckets = new SizeConfigurationBuckets(sizeConfigs);
+
+        // Change screen layout long from not long to long and check that screen layout flag is
+        // still part of the diff because a threshold was crossed.
+        final int diff = CONFIG_SCREEN_LAYOUT;
+        final Configuration oldConfig = new Configuration();
+        oldConfig.screenLayout |= SCREENLAYOUT_LONG_NO;
+        final Configuration newConfig = new Configuration();
+        newConfig.screenLayout |= SCREENLAYOUT_LONG_YES;
+        final int filteredDiff = SizeConfigurationBuckets.filterDiff(diff, oldConfig, newConfig,
+                sizeBuckets);
+        assertEquals(CONFIG_SCREEN_LAYOUT, filteredDiff);
+    }
+
+    /**
+     * Tests that horizontal buckets are correctly checked in
+     * {@link SizeConfigurationBuckets#filterDiff()}.
+     */
+    @Test
+    public void testHorizontalSizeThresholds() {
+        final int[] horizontalThresholds = new int[] { 360, 600 };
+        final SizeConfigurationBuckets sizeBuckets = new SizeConfigurationBuckets(
+                horizontalThresholds, null /* vertical */, null /* smallest */,
+                null /* screenLayoutSize */, false /* screenLayoutLongSet */);
+
+        final Configuration oldConfig = new Configuration();
+        final Configuration newConfig = new Configuration();
+
+        oldConfig.screenWidthDp = 480;
+        // Test that value within bucket filters out screen size config
+        newConfig.screenWidthDp = 520;
+        assertEquals(0, SizeConfigurationBuckets.filterDiff(CONFIG_SCREEN_SIZE, oldConfig,
+                newConfig, sizeBuckets));
+        // Test that value outside bucket does not filter out screen size config
+        newConfig.screenWidthDp = 640;
+        assertEquals(CONFIG_SCREEN_SIZE, SizeConfigurationBuckets.filterDiff(CONFIG_SCREEN_SIZE,
+                oldConfig, newConfig, sizeBuckets));
+    }
+
+    /**
+     * Tests that vertical buckets are correctly checked in
+     * {@link SizeConfigurationBuckets#filterDiff()}.
+     */
+    @Test
+    public void testVerticalSizeThresholds() {
+        final int[] verticalThresholds = new int[] { 360, 600 };
+        final SizeConfigurationBuckets sizeBuckets = new SizeConfigurationBuckets(
+                null, verticalThresholds /* vertical */, null /* smallest */,
+                null /* screenLayoutSize */, false /* screenLayoutLongSet */);
+
+        final Configuration oldConfig = new Configuration();
+        final Configuration newConfig = new Configuration();
+
+        oldConfig.screenHeightDp = 480;
+        // Test that value within bucket filters out screen size config
+        newConfig.screenHeightDp = 520;
+        assertEquals(0, SizeConfigurationBuckets.filterDiff(CONFIG_SCREEN_SIZE, oldConfig,
+                newConfig, sizeBuckets));
+        // Test that value outside bucket does not filter out screen size config
+        newConfig.screenHeightDp = 640;
+        assertEquals(CONFIG_SCREEN_SIZE, SizeConfigurationBuckets.filterDiff(CONFIG_SCREEN_SIZE,
+                oldConfig, newConfig, sizeBuckets));
+    }
+
+    /**
+     * Tests that smallest width buckets are correctly checked in
+     * {@link SizeConfigurationBuckets#filterDiff()}.
+     */
+    @Test
+    public void testSmallestWidthSizeThresholds() {
+        final int[] smallestWidthThresholds = new int[] { 360, 600 };
+        final SizeConfigurationBuckets sizeBuckets = new SizeConfigurationBuckets(
+                null, null /* vertical */, smallestWidthThresholds /* smallest */,
+                null /* screenLayoutSize */, false /* screenLayoutLongSet */);
+
+        final Configuration oldConfig = new Configuration();
+        final Configuration newConfig = new Configuration();
+
+        oldConfig.smallestScreenWidthDp = 480;
+        // Test that value within bucket filters out smallest screen size config
+        newConfig.smallestScreenWidthDp = 520;
+        assertEquals(0, SizeConfigurationBuckets.filterDiff(CONFIG_SMALLEST_SCREEN_SIZE, oldConfig,
+                newConfig, sizeBuckets));
+        // Test that value outside bucket does not filter out smallest screen size config
+        newConfig.smallestScreenWidthDp = 640;
+        assertEquals(CONFIG_SMALLEST_SCREEN_SIZE, SizeConfigurationBuckets.filterDiff(
+                CONFIG_SMALLEST_SCREEN_SIZE, oldConfig, newConfig, sizeBuckets));
+    }
+
+    private void checkCrossesSizeThreshold(int[] thresholds, int oldValue, int newValue,
+            boolean expected) {
+        final String errorString = String.format(
+                "thresholds=%s, oldValue=%d, newValue=%d, expected=%b", Arrays.toString(thresholds),
+                oldValue, newValue, expected);
+        final boolean actual = SizeConfigurationBuckets.crossesSizeThreshold(thresholds, oldValue,
+                newValue);
+        assertEquals(errorString, expected, actual);
+    }
+}