diff --git a/libs/gui/tests/BLASTBufferQueue_test.cpp b/libs/gui/tests/BLASTBufferQueue_test.cpp
index c4c2fa5..cf2593d 100644
--- a/libs/gui/tests/BLASTBufferQueue_test.cpp
+++ b/libs/gui/tests/BLASTBufferQueue_test.cpp
@@ -188,7 +188,10 @@
     void SetUp() {
         mComposer = ComposerService::getComposerService();
         mClient = new SurfaceComposerClient();
-        mDisplayToken = mClient->getInternalDisplayToken();
+        const auto ids = SurfaceComposerClient::getPhysicalDisplayIds();
+        ASSERT_FALSE(ids.empty());
+        // display 0 is picked as this test is not much display depedent
+        mDisplayToken = SurfaceComposerClient::getPhysicalDisplayToken(ids.front());
         ASSERT_NE(nullptr, mDisplayToken.get());
         Transaction t;
         t.setDisplayLayerStack(mDisplayToken, ui::DEFAULT_LAYER_STACK);
diff --git a/libs/gui/tests/DisplayedContentSampling_test.cpp b/libs/gui/tests/DisplayedContentSampling_test.cpp
index b647aab..0a2750a 100644
--- a/libs/gui/tests/DisplayedContentSampling_test.cpp
+++ b/libs/gui/tests/DisplayedContentSampling_test.cpp
@@ -32,7 +32,10 @@
     void SetUp() {
         mComposerClient = new SurfaceComposerClient;
         ASSERT_EQ(OK, mComposerClient->initCheck());
-        mDisplayToken = mComposerClient->getInternalDisplayToken();
+        const auto ids = SurfaceComposerClient::getPhysicalDisplayIds();
+        ASSERT_FALSE(ids.empty());
+        // display 0 is picked for now, can extend to support all displays if needed
+        mDisplayToken = SurfaceComposerClient::getPhysicalDisplayToken(ids.front());
         ASSERT_TRUE(mDisplayToken);
     }
 
diff --git a/libs/gui/tests/EndToEndNativeInputTest.cpp b/libs/gui/tests/EndToEndNativeInputTest.cpp
index 2637f59..3344e0b 100644
--- a/libs/gui/tests/EndToEndNativeInputTest.cpp
+++ b/libs/gui/tests/EndToEndNativeInputTest.cpp
@@ -360,8 +360,10 @@
     void SetUp() {
         mComposerClient = new SurfaceComposerClient;
         ASSERT_EQ(NO_ERROR, mComposerClient->initCheck());
-
-        const auto display = mComposerClient->getInternalDisplayToken();
+        const auto ids = SurfaceComposerClient::getPhysicalDisplayIds();
+        ASSERT_FALSE(ids.empty());
+        // display 0 is picked for now, can extend to support all displays if needed
+        const auto display = SurfaceComposerClient::getPhysicalDisplayToken(ids.front());
         ASSERT_NE(display, nullptr);
 
         ui::DisplayMode mode;
diff --git a/libs/gui/tests/Surface_test.cpp b/libs/gui/tests/Surface_test.cpp
index b9358e7..c078378 100644
--- a/libs/gui/tests/Surface_test.cpp
+++ b/libs/gui/tests/Surface_test.cpp
@@ -263,7 +263,10 @@
     sp<ANativeWindow> anw(mSurface);
 
     // Verify the screenshot works with no protected buffers.
-    const sp<IBinder> display = ComposerServiceAIDL::getInstance().getInternalDisplayToken();
+    const auto ids = SurfaceComposerClient::getPhysicalDisplayIds();
+    ASSERT_FALSE(ids.empty());
+    // display 0 is picked for now, can extend to support all displays if needed
+    const sp<IBinder> display = SurfaceComposerClient::getPhysicalDisplayToken(ids.front());
     ASSERT_FALSE(display == nullptr);
 
     DisplayCaptureArgs captureArgs;