1 files changed, 5 insertions, 0 deletions
diff --git a/services/surfaceflinger/tests/unittests/SurfaceFlinger_SetPowerModeInternalTest.cpp b/services/surfaceflinger/tests/unittests/SurfaceFlinger_SetPowerModeInternalTest.cpp
index 65024202b8..eea1002236 100644
--- a/services/surfaceflinger/tests/unittests/SurfaceFlinger_SetPowerModeInternalTest.cpp
+++ b/services/surfaceflinger/tests/unittests/SurfaceFlinger_SetPowerModeInternalTest.cpp
@@ -260,6 +260,11 @@ struct DisplayPowerCase {
         auto display = Display::makeFakeExistingDisplayInjector(test);
         display.inject();
         display.mutableDisplayDevice()->setPowerMode(mode);
+        if (display.mutableDisplayDevice()->isInternal()) {
+            test->mFlinger.mutableActiveDisplayToken() =
+                    display.mutableDisplayDevice()->getDisplayToken();
+        }
+
         return display;
     }