diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index c5b270ea971a5..2398661668b1b 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2291,9 +2291,10 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
             }
             mPowerManagerWrapper.wakeUp(SystemClock.uptimeMillis(),
                     "android.server.wm:TURN_ON");
-        }
-        if (mAppToken != null) {
-            mAppToken.setCanTurnScreenOn(false);
+
+            if (mAppToken != null) {
+                mAppToken.setCanTurnScreenOn(false);
+            }
         }
 
         // If we were already visible, skip rest of preparation.
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowStateTests.java b/services/tests/servicestests/src/com/android/server/wm/WindowStateTests.java
index 83868d6f20499..678aa642d8e12 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowStateTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowStateTests.java
@@ -27,6 +27,8 @@ import android.support.test.runner.AndroidJUnit4;
 
 import java.util.LinkedList;
 
+import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
+import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.view.WindowManager.LayoutParams.FIRST_SUB_WINDOW;
 import static android.view.WindowManager.LayoutParams.FLAG_ALT_FOCUSABLE_IM;
 import static android.view.WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
@@ -36,6 +38,7 @@ import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_ATTACHED_
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_MEDIA;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_MEDIA_OVERLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_SUB_PANEL;
+import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
 
 import static org.junit.Assert.assertEquals;
@@ -44,6 +47,7 @@ import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.verify;
 
@@ -218,6 +222,60 @@ public class WindowStateTests extends WindowTestsBase {
     public void testPrepareWindowToDisplayDuringRelayout() throws Exception {
         testPrepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
         testPrepareWindowToDisplayDuringRelayout(true /*wasVisible*/);
+
+        // Call prepareWindowToDisplayDuringRelayout for a window without FLAG_TURN_SCREEN_ON
+        // before calling prepareWindowToDisplayDuringRelayout for windows with flag in the same
+        // appWindowToken.
+        final AppWindowToken appWindowToken = createAppWindowToken(mDisplayContent,
+                WINDOWING_MODE_FULLSCREEN, ACTIVITY_TYPE_STANDARD);
+        final WindowState first = createWindow(null, TYPE_APPLICATION, appWindowToken, "first");
+        final WindowState second = createWindow(null, TYPE_APPLICATION, appWindowToken, "second");
+        second.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
+
+        reset(mPowerManagerWrapper);
+        first.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
+        verify(mPowerManagerWrapper, never()).wakeUp(anyLong(), anyString());
+        assertTrue(appWindowToken.canTurnScreenOn());
+
+        reset(mPowerManagerWrapper);
+        second.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
+        verify(mPowerManagerWrapper).wakeUp(anyLong(), anyString());
+        assertFalse(appWindowToken.canTurnScreenOn());
+
+        // Call prepareWindowToDisplayDuringRelayout for two window that have FLAG_TURN_SCREEN_ON
+        // from the same appWindowToken. Only one should trigger the wakeup.
+        appWindowToken.setCanTurnScreenOn(true);
+        first.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
+        second.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
+
+        reset(mPowerManagerWrapper);
+        first.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
+        verify(mPowerManagerWrapper).wakeUp(anyLong(), anyString());
+        assertFalse(appWindowToken.canTurnScreenOn());
+
+        reset(mPowerManagerWrapper);
+        second.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
+        verify(mPowerManagerWrapper, never()).wakeUp(anyLong(), anyString());
+        assertFalse(appWindowToken.canTurnScreenOn());
+
+        // Call prepareWindowToDisplayDuringRelayout for a windows that are not children of an
+        // appWindowToken. Both windows have the FLAG_TURNS_SCREEN_ON so both should call wakeup
+        final WindowToken windowToken = new WindowTestUtils.TestWindowToken(FIRST_SUB_WINDOW,
+                mDisplayContent);
+        final WindowState firstWindow = createWindow(null, TYPE_APPLICATION, windowToken,
+                "firstWindow");
+        final WindowState secondWindow = createWindow(null, TYPE_APPLICATION, windowToken,
+                "secondWindow");
+        firstWindow.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
+        secondWindow.mAttrs.flags |= WindowManager.LayoutParams.FLAG_TURN_SCREEN_ON;
+
+        reset(mPowerManagerWrapper);
+        firstWindow.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
+        verify(mPowerManagerWrapper).wakeUp(anyLong(), anyString());
+
+        reset(mPowerManagerWrapper);
+        secondWindow.prepareWindowToDisplayDuringRelayout(false /*wasVisible*/);
+        verify(mPowerManagerWrapper).wakeUp(anyLong(), anyString());
     }
 
     @Test
diff --git a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
index 91d5ea463bbed..57c1a008e3efa 100644
--- a/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/servicestests/src/com/android/server/wm/WindowTestsBase.java
@@ -198,12 +198,17 @@ class WindowTestsBase {
             return new WindowTestUtils.TestWindowToken(type, dc);
         }
 
-        final TaskStack stack =
-                createStackControllerOnStackOnDisplay(windowingMode, activityType, dc).mContainer;
+        return createAppWindowToken(dc, windowingMode, activityType);
+    }
+
+    AppWindowToken createAppWindowToken(DisplayContent dc, int windowingMode, int activityType) {
+        final TaskStack stack = createStackControllerOnStackOnDisplay(windowingMode, activityType,
+                dc).mContainer;
         final Task task = createTaskInStack(stack, 0 /* userId */);
-        final WindowTestUtils.TestAppWindowToken token = new WindowTestUtils.TestAppWindowToken(dc);
-        task.addChild(token, 0);
-        return token;
+        final WindowTestUtils.TestAppWindowToken appWindowToken =
+                new WindowTestUtils.TestAppWindowToken(mDisplayContent);
+        task.addChild(appWindowToken, 0);
+        return appWindowToken;
     }
 
     WindowState createWindow(WindowState parent, int type, String name) {