diff --git a/services/surfaceflinger/tests/unittests/CompositionTest.cpp b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
index 3042450..560f139 100644
--- a/services/surfaceflinger/tests/unittests/CompositionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
@@ -105,7 +105,9 @@
 
         mFlinger.setupRenderEngine(std::unique_ptr<renderengine::RenderEngine>(mRenderEngine));
         mFlinger.setupTimeStats(std::shared_ptr<TimeStats>(mTimeStats));
-        setupComposer(0);
+
+        mComposer = new Hwc2::mock::Composer();
+        mFlinger.setupComposer(std::unique_ptr<Hwc2::Composer>(mComposer));
     }
 
     ~CompositionTest() {
@@ -114,14 +116,6 @@
         ALOGD("**** Tearing down after %s.%s\n", test_info->test_case_name(), test_info->name());
     }
 
-    void setupComposer(int virtualDisplayCount) {
-        mComposer = new Hwc2::mock::Composer();
-        EXPECT_CALL(*mComposer, getMaxVirtualDisplayCount()).WillOnce(Return(virtualDisplayCount));
-        mFlinger.setupComposer(std::unique_ptr<Hwc2::Composer>(mComposer));
-
-        Mock::VerifyAndClear(mComposer);
-    }
-
     void setupScheduler() {
         auto eventThread = std::make_unique<mock::EventThread>();
         auto sfEventThread = std::make_unique<mock::EventThread>();
@@ -289,16 +283,16 @@
         const ::testing::TestInfo* const test_info =
                 ::testing::UnitTest::GetInstance()->current_test_info();
 
-        auto ceDisplayArgs =
-                compositionengine::DisplayCreationArgsBuilder()
-                        .setPhysical({DEFAULT_DISPLAY_ID, ui::DisplayConnectionType::Internal})
-                        .setPixels({DEFAULT_DISPLAY_WIDTH, DEFAULT_DISPLAY_HEIGHT})
-                        .setIsSecure(Derived::IS_SECURE)
-                        .setLayerStackId(DEFAULT_LAYER_STACK)
-                        .setPowerAdvisor(&test->mPowerAdvisor)
-                        .setName(std::string("Injected display for ") +
-                                 test_info->test_case_name() + "." + test_info->name())
-                        .build();
+        auto ceDisplayArgs = compositionengine::DisplayCreationArgsBuilder()
+                                     .setId(DEFAULT_DISPLAY_ID)
+                                     .setConnectionType(ui::DisplayConnectionType::Internal)
+                                     .setPixels({DEFAULT_DISPLAY_WIDTH, DEFAULT_DISPLAY_HEIGHT})
+                                     .setIsSecure(Derived::IS_SECURE)
+                                     .setLayerStackId(DEFAULT_LAYER_STACK)
+                                     .setPowerAdvisor(&test->mPowerAdvisor)
+                                     .setName(std::string("Injected display for ") +
+                                              test_info->test_case_name() + "." + test_info->name())
+                                     .build();
 
         auto compositionDisplay =
                 compositionengine::impl::createDisplay(test->mFlinger.getCompositionEngine(),
diff --git a/services/surfaceflinger/tests/unittests/DisplayIdGeneratorTest.cpp b/services/surfaceflinger/tests/unittests/DisplayIdGeneratorTest.cpp
index 77a3e14..8d4a023 100644
--- a/services/surfaceflinger/tests/unittests/DisplayIdGeneratorTest.cpp
+++ b/services/surfaceflinger/tests/unittests/DisplayIdGeneratorTest.cpp
@@ -14,76 +14,68 @@
  * limitations under the License.
  */
 
-// TODO(b/129481165): remove the #pragma below and fix conversion issues
-#pragma clang diagnostic push
-#pragma clang diagnostic ignored "-Wextra"
-
 #include <gtest/gtest.h>
+#include <ui/DisplayId.h>
 
+#include <algorithm>
+#include <iterator>
 #include <vector>
 
-#include <ui/DisplayId.h>
 #include "DisplayIdGenerator.h"
 
 namespace android {
 
-template <typename T>
-void testNextId(DisplayIdGenerator<T>& generator) {
-    constexpr int kNumIds = 5;
-    std::vector<T> ids;
-    for (int i = 0; i < kNumIds; i++) {
-        const auto id = generator.nextId();
-        ASSERT_TRUE(id);
-        ids.push_back(*id);
-    }
+template <typename Id>
+void testGenerateId() {
+    DisplayIdGenerator<Id> generator;
+
+    std::vector<std::optional<Id>> ids;
+    std::generate_n(std::back_inserter(ids), 10, [&] { return generator.generateId(); });
 
     // All IDs should be different.
-    for (size_t i = 0; i < kNumIds; i++) {
-        for (size_t j = i + 1; j < kNumIds; j++) {
-            EXPECT_NE(ids[i], ids[j]);
+    for (auto it = ids.begin(); it != ids.end(); ++it) {
+        EXPECT_TRUE(*it);
+
+        for (auto dup = it + 1; dup != ids.end(); ++dup) {
+            EXPECT_NE(*it, *dup);
         }
     }
 }
 
-TEST(DisplayIdGeneratorTest, nextIdGpuVirtual) {
-    RandomDisplayIdGenerator<GpuVirtualDisplayId> generator;
-    testNextId(generator);
+TEST(DisplayIdGeneratorTest, generateGpuVirtualDisplayId) {
+    testGenerateId<GpuVirtualDisplayId>();
 }
 
-TEST(DisplayIdGeneratorTest, nextIdHalVirtual) {
-    RandomDisplayIdGenerator<HalVirtualDisplayId> generator;
-    testNextId(generator);
+TEST(DisplayIdGeneratorTest, generateHalVirtualDisplayId) {
+    testGenerateId<HalVirtualDisplayId>();
 }
 
-TEST(DisplayIdGeneratorTest, markUnused) {
+TEST(DisplayIdGeneratorTest, releaseId) {
     constexpr size_t kMaxIdsCount = 5;
-    RandomDisplayIdGenerator<GpuVirtualDisplayId> generator(kMaxIdsCount);
+    DisplayIdGenerator<GpuVirtualDisplayId> generator(kMaxIdsCount);
 
-    const auto id = generator.nextId();
+    const auto id = generator.generateId();
     EXPECT_TRUE(id);
 
-    for (int i = 1; i < kMaxIdsCount; i++) {
-        EXPECT_TRUE(generator.nextId());
+    for (size_t i = 1; i < kMaxIdsCount; i++) {
+        EXPECT_TRUE(generator.generateId());
     }
 
-    EXPECT_FALSE(generator.nextId());
+    EXPECT_FALSE(generator.generateId());
 
-    generator.markUnused(*id);
-    EXPECT_TRUE(generator.nextId());
+    generator.releaseId(*id);
+    EXPECT_TRUE(generator.generateId());
 }
 
 TEST(DisplayIdGeneratorTest, maxIdsCount) {
     constexpr size_t kMaxIdsCount = 5;
-    RandomDisplayIdGenerator<GpuVirtualDisplayId> generator(kMaxIdsCount);
+    DisplayIdGenerator<GpuVirtualDisplayId> generator(kMaxIdsCount);
 
-    for (int i = 0; i < kMaxIdsCount; i++) {
-        EXPECT_TRUE(generator.nextId());
+    for (size_t i = 0; i < kMaxIdsCount; i++) {
+        EXPECT_TRUE(generator.generateId());
     }
 
-    EXPECT_FALSE(generator.nextId());
+    EXPECT_FALSE(generator.generateId());
 }
 
 } // namespace android
-
-// TODO(b/129481165): remove the #pragma below and fix conversion issues
-#pragma clang diagnostic pop // ignored "-Wextra"
\ No newline at end of file
diff --git a/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp b/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
index a3e8108..cc24323 100644
--- a/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
@@ -41,9 +41,6 @@
     mFlinger.mutableUseColorManagement() = false;
     mFlinger.mutableDisplayColorSetting() = DisplayColorSetting::kUnmanaged;
 
-    // Default to using HWC virtual displays
-    mFlinger.mutableUseHwcVirtualDisplays() = true;
-
     mFlinger.setCreateBufferQueueFunction([](auto, auto, auto) {
         ADD_FAILURE() << "Unexpected request to create a buffer queue.";
     });
@@ -85,10 +82,17 @@
 }
 
 void DisplayTransactionTest::injectMockComposer(int virtualDisplayCount) {
+    if (mComposer) {
+        // If reinjecting, disable first to prevent the enable below from being a no-op.
+        mFlinger.enableHalVirtualDisplays(false);
+    }
+
     mComposer = new Hwc2::mock::Composer();
-    EXPECT_CALL(*mComposer, getMaxVirtualDisplayCount()).WillOnce(Return(virtualDisplayCount));
     mFlinger.setupComposer(std::unique_ptr<Hwc2::Composer>(mComposer));
 
+    EXPECT_CALL(*mComposer, getMaxVirtualDisplayCount()).WillOnce(Return(virtualDisplayCount));
+    mFlinger.enableHalVirtualDisplays(true);
+
     Mock::VerifyAndClear(mComposer);
 }
 
@@ -135,18 +139,21 @@
     EXPECT_CALL(*mNativeWindow, perform(NATIVE_WINDOW_SET_USAGE64));
     EXPECT_CALL(*mNativeWindow, perform(NATIVE_WINDOW_API_DISCONNECT)).Times(AnyNumber());
 
-    auto compositionDisplay = compositionengine::impl::
-            createDisplay(mFlinger.getCompositionEngine(),
-                          compositionengine::DisplayCreationArgsBuilder()
-                                  .setPhysical(
-                                          {DEFAULT_DISPLAY_ID, ui::DisplayConnectionType::Internal})
-                                  .setPixels({DEFAULT_DISPLAY_WIDTH, DEFAULT_DISPLAY_HEIGHT})
-                                  .setPowerAdvisor(&mPowerAdvisor)
-                                  .build());
+    constexpr auto kConnectionType = ui::DisplayConnectionType::Internal;
+    constexpr bool kIsPrimary = true;
 
-    auto injector = FakeDisplayDeviceInjector(mFlinger, compositionDisplay,
-                                              ui::DisplayConnectionType::Internal,
-                                              DEFAULT_DISPLAY_HWC_DISPLAY_ID, true /* isPrimary */);
+    auto compositionDisplay =
+            compositionengine::impl::createDisplay(mFlinger.getCompositionEngine(),
+                                                   compositionengine::DisplayCreationArgsBuilder()
+                                                           .setId(DEFAULT_DISPLAY_ID)
+                                                           .setConnectionType(kConnectionType)
+                                                           .setPixels({DEFAULT_DISPLAY_WIDTH,
+                                                                       DEFAULT_DISPLAY_HEIGHT})
+                                                           .setPowerAdvisor(&mPowerAdvisor)
+                                                           .build());
+
+    auto injector = FakeDisplayDeviceInjector(mFlinger, compositionDisplay, kConnectionType,
+                                              DEFAULT_DISPLAY_HWC_DISPLAY_ID, kIsPrimary);
 
     injector.setNativeWindow(mNativeWindow);
     if (injectExtra) {
diff --git a/services/surfaceflinger/tests/unittests/DisplayTransactionTestHelpers.h b/services/surfaceflinger/tests/unittests/DisplayTransactionTestHelpers.h
index d68fff6..6ce281d 100644
--- a/services/surfaceflinger/tests/unittests/DisplayTransactionTestHelpers.h
+++ b/services/surfaceflinger/tests/unittests/DisplayTransactionTestHelpers.h
@@ -263,40 +263,23 @@
 
     static auto makeFakeExistingDisplayInjector(DisplayTransactionTest* test) {
         auto ceDisplayArgs = compositionengine::DisplayCreationArgsBuilder();
-        if (auto displayId = PhysicalDisplayId::tryCast(DISPLAY_ID::get())) {
-            ceDisplayArgs.setPhysical({*displayId, ui::DisplayConnectionType::Internal});
-        } else {
-            // We turn off the use of HwcVirtualDisplays, to prevent Composition Engine
-            // from calling into HWComposer. This way all virtual displays will get
-            // a GpuVirtualDisplayId, even if we are in the HwcVirtualDisplayVariant.
-            // In this case we later override it by calling display.setDisplayIdForTesting().
-            ceDisplayArgs.setUseHwcVirtualDisplays(false);
+        ceDisplayArgs.setId(DISPLAY_ID::get())
+                .setPixels({WIDTH, HEIGHT})
+                .setPowerAdvisor(&test->mPowerAdvisor);
 
-            GpuVirtualDisplayId desiredDisplayId = GpuVirtualDisplayId::tryCast(DISPLAY_ID::get())
-                                                           .value_or(GpuVirtualDisplayId(0));
-
-            ON_CALL(test->mFlinger.gpuVirtualDisplayIdGenerator(), nextId())
-                    .WillByDefault(Return(desiredDisplayId));
-
-            auto& generator = test->mFlinger.gpuVirtualDisplayIdGenerator();
-            ceDisplayArgs.setGpuVirtualDisplayIdGenerator(generator);
+        const auto connectionType = CONNECTION_TYPE::value;
+        if (connectionType) {
+            ceDisplayArgs.setConnectionType(*connectionType);
         }
-        ceDisplayArgs.setPixels({WIDTH, HEIGHT}).setPowerAdvisor(&test->mPowerAdvisor);
 
         auto compositionDisplay =
                 compositionengine::impl::createDisplay(test->mFlinger.getCompositionEngine(),
                                                        ceDisplayArgs.build());
 
-        if (HalVirtualDisplayId::tryCast(DISPLAY_ID::get())) {
-            // CompositionEngine has assigned a placeholder GpuVirtualDisplayId and we need to
-            // override it with the correct HalVirtualDisplayId.
-            compositionDisplay->setDisplayIdForTesting(DISPLAY_ID::get());
-        }
-
         auto injector =
                 TestableSurfaceFlinger::FakeDisplayDeviceInjector(test->mFlinger,
                                                                   compositionDisplay,
-                                                                  CONNECTION_TYPE::value,
+                                                                  connectionType,
                                                                   HWC_DISPLAY_ID_OPT::value,
                                                                   static_cast<bool>(PRIMARY));
 
@@ -404,8 +387,8 @@
                 ::testing::UnitTest::GetInstance()->current_test_info();
 
         auto ceDisplayArgs = compositionengine::DisplayCreationArgsBuilder()
-                                     .setPhysical({DisplayVariant::DISPLAY_ID::get(),
-                                                   PhysicalDisplay::CONNECTION_TYPE})
+                                     .setId(DisplayVariant::DISPLAY_ID::get())
+                                     .setConnectionType(PhysicalDisplay::CONNECTION_TYPE)
                                      .setPixels({DisplayVariant::WIDTH, DisplayVariant::HEIGHT})
                                      .setIsSecure(static_cast<bool>(DisplayVariant::SECURE))
                                      .setPowerAdvisor(&test->mPowerAdvisor)
@@ -558,17 +541,13 @@
         const ::testing::TestInfo* const test_info =
                 ::testing::UnitTest::GetInstance()->current_test_info();
 
-        ON_CALL(test->mFlinger.gpuVirtualDisplayIdGenerator(), nextId())
-                .WillByDefault(Return(Base::DISPLAY_ID::get()));
-
         auto ceDisplayArgs = compositionengine::DisplayCreationArgsBuilder()
+                                     .setId(Base::DISPLAY_ID::get())
                                      .setPixels({Base::WIDTH, Base::HEIGHT})
                                      .setIsSecure(static_cast<bool>(Base::SECURE))
                                      .setPowerAdvisor(&test->mPowerAdvisor)
                                      .setName(std::string("Injected display for ") +
                                               test_info->test_case_name() + "." + test_info->name())
-                                     .setGpuVirtualDisplayIdGenerator(
-                                             test->mFlinger.gpuVirtualDisplayIdGenerator())
                                      .build();
 
         return compositionengine::impl::createDisplay(test->mFlinger.getCompositionEngine(),
@@ -610,35 +589,22 @@
         const ::testing::TestInfo* const test_info =
                 ::testing::UnitTest::GetInstance()->current_test_info();
 
-        // In order to prevent compostition engine calling into HWComposer, we
-        // 1. turn off the use of HWC virtual displays,
-        // 2. provide a GpuVirtualDisplayIdGenerator which always returns some fake ID
-        // 3. override the ID by calling setDisplayIdForTesting()
-
-        ON_CALL(test->mFlinger.gpuVirtualDisplayIdGenerator(), nextId())
-                .WillByDefault(Return(GpuVirtualDisplayId(0)));
-
+        const auto displayId = Base::DISPLAY_ID::get();
         auto ceDisplayArgs = compositionengine::DisplayCreationArgsBuilder()
-                                     .setUseHwcVirtualDisplays(false)
+                                     .setId(displayId)
                                      .setPixels({Base::WIDTH, Base::HEIGHT})
                                      .setIsSecure(static_cast<bool>(Base::SECURE))
                                      .setPowerAdvisor(&test->mPowerAdvisor)
                                      .setName(std::string("Injected display for ") +
                                               test_info->test_case_name() + "." + test_info->name())
-                                     .setGpuVirtualDisplayIdGenerator(
-                                             test->mFlinger.gpuVirtualDisplayIdGenerator())
                                      .build();
 
         auto compositionDisplay =
                 compositionengine::impl::createDisplay(test->mFlinger.getCompositionEngine(),
                                                        ceDisplayArgs);
-        compositionDisplay->setDisplayIdForTesting(Base::DISPLAY_ID::get());
 
         // Insert display data so that the HWC thinks it created the virtual display.
-        if (const auto displayId = Base::DISPLAY_ID::get();
-            HalVirtualDisplayId::tryCast(displayId)) {
-            test->mFlinger.mutableHwcDisplayData().try_emplace(displayId);
-        }
+        test->mFlinger.mutableHwcDisplayData().try_emplace(displayId);
 
         return compositionDisplay;
     }
@@ -649,8 +615,8 @@
     }
 
     static void setupHwcVirtualDisplayCreationCallExpectations(DisplayTransactionTest* test) {
-        EXPECT_CALL(*test->mComposer, createVirtualDisplay(Base::WIDTH, Base::HEIGHT, _, _))
-                .WillOnce(DoAll(SetArgPointee<3>(Self::HWC_DISPLAY_ID), Return(Error::NONE)));
+        EXPECT_CALL(*test->mComposer, createVirtualDisplay(Base::WIDTH, Base::HEIGHT, _, _, _))
+                .WillOnce(DoAll(SetArgPointee<4>(Self::HWC_DISPLAY_ID), Return(Error::NONE)));
         EXPECT_CALL(*test->mComposer, setClientTargetSlotCount(_)).WillOnce(Return(Error::NONE));
     }
 };
diff --git a/services/surfaceflinger/tests/unittests/FpsReporterTest.cpp b/services/surfaceflinger/tests/unittests/FpsReporterTest.cpp
index dec0ff5..010c675 100644
--- a/services/surfaceflinger/tests/unittests/FpsReporterTest.cpp
+++ b/services/surfaceflinger/tests/unittests/FpsReporterTest.cpp
@@ -76,11 +76,9 @@
     static constexpr int32_t PRIORITY_UNSET = -1;
 
     void setupScheduler();
-    void setupComposer(uint32_t virtualDisplayCount);
     sp<BufferStateLayer> createBufferStateLayer(LayerMetadata metadata);
 
     TestableSurfaceFlinger mFlinger;
-    Hwc2::mock::Composer* mComposer = nullptr;
     mock::FrameTimeline mFrameTimeline =
             mock::FrameTimeline(std::make_shared<impl::TimeStats>(), 0);
 
@@ -103,7 +101,8 @@
     ALOGD("**** Setting up for %s.%s\n", test_info->test_case_name(), test_info->name());
 
     setupScheduler();
-    setupComposer(0);
+    mFlinger.setupComposer(std::make_unique<Hwc2::mock::Composer>());
+
     mFpsListener = new TestableFpsListener();
 }
 
@@ -145,19 +144,11 @@
                             std::move(eventThread), std::move(sfEventThread));
 }
 
-void FpsReporterTest::setupComposer(uint32_t virtualDisplayCount) {
-    mComposer = new Hwc2::mock::Composer();
-    EXPECT_CALL(*mComposer, getMaxVirtualDisplayCount()).WillOnce(Return(virtualDisplayCount));
-    mFlinger.setupComposer(std::unique_ptr<Hwc2::Composer>(mComposer));
-
-    Mock::VerifyAndClear(mComposer);
-}
-
 namespace {
 
 TEST_F(FpsReporterTest, callsListeners) {
     mParent = createBufferStateLayer();
-    const constexpr int32_t kTaskId = 12;
+    constexpr int32_t kTaskId = 12;
     LayerMetadata targetMetadata;
     targetMetadata.setInt32(METADATA_TASK_ID, kTaskId);
     mTarget = createBufferStateLayer(targetMetadata);
diff --git a/services/surfaceflinger/tests/unittests/HWComposerTest.cpp b/services/surfaceflinger/tests/unittests/HWComposerTest.cpp
index cbf8cc2..57f065a 100644
--- a/services/surfaceflinger/tests/unittests/HWComposerTest.cpp
+++ b/services/surfaceflinger/tests/unittests/HWComposerTest.cpp
@@ -79,7 +79,6 @@
     const std::string kMetadata2Name = "com.example.metadata.2";
     constexpr bool kMetadata2Mandatory = true;
 
-    EXPECT_CALL(*mHal, getMaxVirtualDisplayCount()).WillOnce(Return(0));
     EXPECT_CALL(*mHal, getCapabilities()).WillOnce(Return(std::vector<hal::Capability>{}));
     EXPECT_CALL(*mHal, getLayerGenericMetadataKeys(_))
             .WillOnce(DoAll(SetArgPointee<0>(std::vector<hal::LayerGenericMetadataKey>{
@@ -102,7 +101,6 @@
 }
 
 TEST_F(HWComposerSetConfigurationTest, handlesUnsupportedCallToGetLayerGenericMetadataKeys) {
-    EXPECT_CALL(*mHal, getMaxVirtualDisplayCount()).WillOnce(Return(0));
     EXPECT_CALL(*mHal, getCapabilities()).WillOnce(Return(std::vector<hal::Capability>{}));
     EXPECT_CALL(*mHal, getLayerGenericMetadataKeys(_))
             .WillOnce(Return(hardware::graphics::composer::V2_4::Error::UNSUPPORTED));
diff --git a/services/surfaceflinger/tests/unittests/RefreshRateSelectionTest.cpp b/services/surfaceflinger/tests/unittests/RefreshRateSelectionTest.cpp
index 9c6ad06..fd3e564 100644
--- a/services/surfaceflinger/tests/unittests/RefreshRateSelectionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/RefreshRateSelectionTest.cpp
@@ -60,7 +60,6 @@
     static constexpr int32_t PRIORITY_UNSET = -1;
 
     void setupScheduler();
-    void setupComposer(uint32_t virtualDisplayCount);
     sp<BufferQueueLayer> createBufferQueueLayer();
     sp<BufferStateLayer> createBufferStateLayer();
     sp<EffectLayer> createEffectLayer();
@@ -69,7 +68,6 @@
     void commitTransaction(Layer* layer);
 
     TestableSurfaceFlinger mFlinger;
-    Hwc2::mock::Composer* mComposer = nullptr;
 
     sp<Client> mClient;
     sp<Layer> mParent;
@@ -83,7 +81,7 @@
     ALOGD("**** Setting up for %s.%s\n", test_info->test_case_name(), test_info->name());
 
     setupScheduler();
-    setupComposer(0);
+    mFlinger.setupComposer(std::make_unique<Hwc2::mock::Composer>());
 }
 
 RefreshRateSelectionTest::~RefreshRateSelectionTest() {
@@ -147,14 +145,6 @@
                             std::move(eventThread), std::move(sfEventThread));
 }
 
-void RefreshRateSelectionTest::setupComposer(uint32_t virtualDisplayCount) {
-    mComposer = new Hwc2::mock::Composer();
-    EXPECT_CALL(*mComposer, getMaxVirtualDisplayCount()).WillOnce(Return(virtualDisplayCount));
-    mFlinger.setupComposer(std::unique_ptr<Hwc2::Composer>(mComposer));
-
-    Mock::VerifyAndClear(mComposer);
-}
-
 namespace {
 /* ------------------------------------------------------------------------
  * Test cases
diff --git a/services/surfaceflinger/tests/unittests/SetFrameRateTest.cpp b/services/surfaceflinger/tests/unittests/SetFrameRateTest.cpp
index c088ddc..46ef750 100644
--- a/services/surfaceflinger/tests/unittests/SetFrameRateTest.cpp
+++ b/services/surfaceflinger/tests/unittests/SetFrameRateTest.cpp
@@ -118,14 +118,12 @@
     SetFrameRateTest();
 
     void setupScheduler();
-    void setupComposer(uint32_t virtualDisplayCount);
 
     void addChild(sp<Layer> layer, sp<Layer> child);
     void removeChild(sp<Layer> layer, sp<Layer> child);
     void commitTransaction();
 
     TestableSurfaceFlinger mFlinger;
-    Hwc2::mock::Composer* mComposer = nullptr;
     mock::MessageQueue* mMessageQueue = new mock::MessageQueue();
 
     std::vector<sp<Layer>> mLayers;
@@ -139,10 +137,11 @@
     mFlinger.mutableUseFrameRateApi() = true;
 
     setupScheduler();
-    setupComposer(0);
 
+    mFlinger.setupComposer(std::make_unique<Hwc2::mock::Composer>());
     mFlinger.mutableEventQueue().reset(mMessageQueue);
 }
+
 void SetFrameRateTest::addChild(sp<Layer> layer, sp<Layer> child) {
     layer.get()->addChild(child.get());
 }
@@ -184,14 +183,6 @@
                             /*hasMultipleModes*/ true);
 }
 
-void SetFrameRateTest::setupComposer(uint32_t virtualDisplayCount) {
-    mComposer = new Hwc2::mock::Composer();
-    EXPECT_CALL(*mComposer, getMaxVirtualDisplayCount()).WillOnce(Return(virtualDisplayCount));
-    mFlinger.setupComposer(std::unique_ptr<Hwc2::Composer>(mComposer));
-
-    Mock::VerifyAndClear(mComposer);
-}
-
 namespace {
 /* ------------------------------------------------------------------------
  * Test cases
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index a551248..9b904d3 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -42,7 +42,6 @@
 #include "SurfaceInterceptor.h"
 #include "TestableScheduler.h"
 #include "mock/DisplayHardware/MockComposer.h"
-#include "mock/MockDisplayIdGenerator.h"
 #include "mock/MockFrameTimeline.h"
 #include "mock/MockFrameTracer.h"
 
@@ -185,9 +184,6 @@
 
     SurfaceFlinger* flinger() { return mFlinger.get(); }
     TestableScheduler* scheduler() { return mScheduler; }
-    mock::DisplayIdGenerator<GpuVirtualDisplayId>& gpuVirtualDisplayIdGenerator() {
-        return mGpuVirtualDisplayIdGenerator;
-    }
 
     // Extend this as needed for accessing SurfaceFlinger private (and public)
     // functions.
@@ -309,6 +305,8 @@
         return mFlinger->destroyDisplay(displayToken);
     }
 
+    void enableHalVirtualDisplays(bool enable) { mFlinger->enableHalVirtualDisplays(enable); }
+
     auto setupNewDisplayDeviceInternal(
             const wp<IBinder>& displayToken,
             std::shared_ptr<compositionengine::Display> compositionDisplay,
@@ -437,7 +435,6 @@
     auto& mutablePhysicalDisplayTokens() { return mFlinger->mPhysicalDisplayTokens; }
     auto& mutableTexturePool() { return mFlinger->mTexturePool; }
     auto& mutableTransactionFlags() { return mFlinger->mTransactionFlags; }
-    auto& mutableUseHwcVirtualDisplays() { return mFlinger->mUseHwcVirtualDisplays; }
     auto& mutablePowerAdvisor() { return mFlinger->mPowerAdvisor; }
     auto& mutableDebugDisableHWC() { return mFlinger->mDebugDisableHWC; }
 
@@ -776,7 +773,6 @@
     surfaceflinger::test::Factory mFactory;
     sp<SurfaceFlinger> mFlinger = new SurfaceFlinger(mFactory, SurfaceFlinger::SkipInitialization);
     TestableScheduler* mScheduler = nullptr;
-    mock::DisplayIdGenerator<GpuVirtualDisplayId> mGpuVirtualDisplayIdGenerator;
 };
 
 } // namespace android
diff --git a/services/surfaceflinger/tests/unittests/TransactionFrameTracerTest.cpp b/services/surfaceflinger/tests/unittests/TransactionFrameTracerTest.cpp
index 25001d3..546bc4a 100644
--- a/services/surfaceflinger/tests/unittests/TransactionFrameTracerTest.cpp
+++ b/services/surfaceflinger/tests/unittests/TransactionFrameTracerTest.cpp
@@ -45,7 +45,7 @@
                 ::testing::UnitTest::GetInstance()->current_test_info();
         ALOGD("**** Setting up for %s.%s\n", test_info->test_case_name(), test_info->name());
         setupScheduler();
-        setupComposer(0);
+        mFlinger.setupComposer(std::make_unique<Hwc2::mock::Composer>());
     }
 
     ~TransactionFrameTracerTest() {
@@ -91,17 +91,9 @@
                                 std::move(eventThread), std::move(sfEventThread));
     }
 
-    void setupComposer(uint32_t virtualDisplayCount) {
-        mComposer = new Hwc2::mock::Composer();
-        EXPECT_CALL(*mComposer, getMaxVirtualDisplayCount()).WillOnce(Return(virtualDisplayCount));
-        mFlinger.setupComposer(std::unique_ptr<Hwc2::Composer>(mComposer));
-
-        Mock::VerifyAndClear(mComposer);
-    }
-
     TestableSurfaceFlinger mFlinger;
-    Hwc2::mock::Composer* mComposer = nullptr;
     renderengine::mock::RenderEngine mRenderEngine;
+
     FenceToFenceTimeMap fenceFactory;
     client_cache_t mClientCache;
 
diff --git a/services/surfaceflinger/tests/unittests/TransactionSurfaceFrameTest.cpp b/services/surfaceflinger/tests/unittests/TransactionSurfaceFrameTest.cpp
index b7917aa..c1123cd 100644
--- a/services/surfaceflinger/tests/unittests/TransactionSurfaceFrameTest.cpp
+++ b/services/surfaceflinger/tests/unittests/TransactionSurfaceFrameTest.cpp
@@ -45,7 +45,7 @@
                 ::testing::UnitTest::GetInstance()->current_test_info();
         ALOGD("**** Setting up for %s.%s\n", test_info->test_case_name(), test_info->name());
         setupScheduler();
-        setupComposer(0);
+        mFlinger.setupComposer(std::make_unique<Hwc2::mock::Composer>());
     }
 
     ~TransactionSurfaceFrameTest() {
@@ -91,17 +91,9 @@
                                 std::move(eventThread), std::move(sfEventThread));
     }
 
-    void setupComposer(uint32_t virtualDisplayCount) {
-        mComposer = new Hwc2::mock::Composer();
-        EXPECT_CALL(*mComposer, getMaxVirtualDisplayCount()).WillOnce(Return(virtualDisplayCount));
-        mFlinger.setupComposer(std::unique_ptr<Hwc2::Composer>(mComposer));
-
-        Mock::VerifyAndClear(mComposer);
-    }
-
     TestableSurfaceFlinger mFlinger;
-    Hwc2::mock::Composer* mComposer = nullptr;
     renderengine::mock::RenderEngine mRenderEngine;
+
     FenceToFenceTimeMap fenceFactory;
     client_cache_t mClientCache;
 
diff --git a/services/surfaceflinger/tests/unittests/mock/DisplayHardware/MockComposer.h b/services/surfaceflinger/tests/unittests/mock/DisplayHardware/MockComposer.h
index 1ba3c0f..cb3bd73 100644
--- a/services/surfaceflinger/tests/unittests/mock/DisplayHardware/MockComposer.h
+++ b/services/surfaceflinger/tests/unittests/mock/DisplayHardware/MockComposer.h
@@ -54,7 +54,8 @@
     MOCK_METHOD0(resetCommands, void());
     MOCK_METHOD0(executeCommands, Error());
     MOCK_METHOD0(getMaxVirtualDisplayCount, uint32_t());
-    MOCK_METHOD4(createVirtualDisplay, Error(uint32_t, uint32_t, PixelFormat*, Display*));
+    MOCK_METHOD5(createVirtualDisplay,
+                 Error(uint32_t, uint32_t, PixelFormat*, std::optional<Display>, Display*));
     MOCK_METHOD1(destroyVirtualDisplay, Error(Display));
     MOCK_METHOD1(acceptDisplayChanges, Error(Display));
     MOCK_METHOD2(createLayer, Error(Display, Layer* outLayer));
diff --git a/services/surfaceflinger/tests/unittests/mock/MockDisplayIdGenerator.h b/services/surfaceflinger/tests/unittests/mock/MockDisplayIdGenerator.h
deleted file mode 100644
index cfc37ea..0000000
--- a/services/surfaceflinger/tests/unittests/mock/MockDisplayIdGenerator.h
+++ /dev/null
@@ -1,36 +0,0 @@
-/*
- * Copyright 2020 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.
- */
-
-#pragma once
-
-#include <gmock/gmock.h>
-
-#include "DisplayIdGenerator.h"
-
-namespace android::mock {
-
-template <typename T>
-class DisplayIdGenerator : public android::DisplayIdGenerator<T> {
-public:
-    // Explicit default instantiation is recommended.
-    DisplayIdGenerator() = default;
-    virtual ~DisplayIdGenerator() = default;
-
-    MOCK_METHOD0(nextId, std::optional<T>());
-    MOCK_METHOD1(markUnused, void(T));
-};
-
-} // namespace android::mock