diff --git a/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp b/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
index 21590df..48b9e92 100644
--- a/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/DisplayTransactionTest.cpp
@@ -44,6 +44,8 @@
 using android::Hwc2::IComposer;
 using android::Hwc2::IComposerClient;
 
+using HWC2Display = TestableSurfaceFlinger::HWC2Display;
+
 constexpr int32_t DEFAULT_REFRESH_RATE = 1666666666;
 constexpr int32_t DEFAULT_DPI = 320;
 
@@ -57,6 +59,8 @@
 
     void expectFramebufferQueuePairCreation(int width, int height);
 
+    std::unordered_set<HWC2::Capability> mCapabilities;
+
     TestableSurfaceFlinger mFlinger;
     mock::EventThread* mEventThread = new mock::EventThread();
 
@@ -69,6 +73,7 @@
     // These mocks are created only when expected to be created via a factory.
     sp<mock::GraphicBufferConsumer> mConsumer;
     sp<mock::GraphicBufferProducer> mProducer;
+    std::unique_ptr<HWC2Display> mDisplay;
 };
 
 DisplayTransactionTest::DisplayTransactionTest() {
@@ -102,36 +107,22 @@
 }
 
 void DisplayTransactionTest::setupPrimaryDisplay(int width, int height) {
-    EXPECT_CALL(*mComposer, getDisplayType(DisplayDevice::DISPLAY_PRIMARY, _))
-            .WillOnce(DoAll(SetArgPointee<1>(IComposerClient::DisplayType::PHYSICAL),
-                            Return(Error::NONE)));
-    EXPECT_CALL(*mComposer, setClientTargetSlotCount(_)).WillOnce(Return(Error::NONE));
-    EXPECT_CALL(*mComposer, getDisplayConfigs(_, _))
-            .WillOnce(DoAll(SetArgPointee<1>(std::vector<unsigned>{0}), Return(Error::NONE)));
-    EXPECT_CALL(*mComposer,
-                getDisplayAttribute(DisplayDevice::DISPLAY_PRIMARY, 0,
-                                    IComposerClient::Attribute::WIDTH, _))
-            .WillOnce(DoAll(SetArgPointee<3>(width), Return(Error::NONE)));
-    EXPECT_CALL(*mComposer,
-                getDisplayAttribute(DisplayDevice::DISPLAY_PRIMARY, 0,
-                                    IComposerClient::Attribute::HEIGHT, _))
-            .WillOnce(DoAll(SetArgPointee<3>(height), Return(Error::NONE)));
-    EXPECT_CALL(*mComposer,
-                getDisplayAttribute(DisplayDevice::DISPLAY_PRIMARY, 0,
-                                    IComposerClient::Attribute::VSYNC_PERIOD, _))
-            .WillOnce(DoAll(SetArgPointee<3>(DEFAULT_REFRESH_RATE), Return(Error::NONE)));
-    EXPECT_CALL(*mComposer,
-                getDisplayAttribute(DisplayDevice::DISPLAY_PRIMARY, 0,
-                                    IComposerClient::Attribute::DPI_X, _))
-            .WillOnce(DoAll(SetArgPointee<3>(DEFAULT_DPI), Return(Error::NONE)));
-    EXPECT_CALL(*mComposer,
-                getDisplayAttribute(DisplayDevice::DISPLAY_PRIMARY, 0,
-                                    IComposerClient::Attribute::DPI_Y, _))
-            .WillOnce(DoAll(SetArgPointee<3>(DEFAULT_DPI), Return(Error::NONE)));
+    mDisplay = std::make_unique<HWC2Display>(*mComposer, mCapabilities, 0,
+                                             HWC2::DisplayType::Physical);
 
-    mFlinger.setupPrimaryDisplay();
+    mDisplay->mutableIsConnected() = true;
+    mDisplay->mutableConfigs().emplace(0,
+                                       HWC2::Display::Config::Builder(*mDisplay, 0)
+                                               .setWidth(width)
+                                               .setHeight(height)
+                                               .setVsyncPeriod(DEFAULT_REFRESH_RATE)
+                                               .setDpiX(DEFAULT_DPI)
+                                               .setDpiY(DEFAULT_DPI)
+                                               .build());
 
-    Mock::VerifyAndClear(mComposer);
+    mFlinger.mutableHwcDisplayData()[DisplayDevice::DISPLAY_PRIMARY].reset();
+    mFlinger.mutableHwcDisplayData()[DisplayDevice::DISPLAY_PRIMARY].hwcDisplay = mDisplay.get();
+    mFlinger.mutableHwcDisplaySlots().emplace(0, DisplayDevice::DISPLAY_PRIMARY);
 }
 
 void DisplayTransactionTest::expectFramebufferQueuePairCreation(int width, int height) {
@@ -196,9 +187,6 @@
     const auto& draw = mFlinger.mutableDrawingState().displays[i];
     EXPECT_EQ(DisplayDevice::DISPLAY_PRIMARY, draw.type);
 
-    EXPECT_CALL(*mComposer, setVsyncEnabled(0, IComposerClient::Vsync::DISABLE))
-            .WillOnce(Return(Error::NONE));
-
     EXPECT_CALL(*mConsumer, consumerDisconnect()).Times(1);
 }
 
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index 4895e16..067b09a 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -44,12 +44,6 @@
         mFlinger->getBE().mHwc.reset(new HWComposer(std::move(composer)));
     }
 
-    void setupPrimaryDisplay() {
-        mFlinger->getBE().mHwc->mHwcDevice->onHotplug(0, HWC2::Connection::Connected);
-        mFlinger->getBE().mHwc->onHotplug(0, DisplayDevice::DISPLAY_PRIMARY,
-                                          HWC2::Connection::Connected);
-    }
-
     using CreateBufferQueueFunction = SurfaceFlinger::CreateBufferQueueFunction;
 
     void setCreateBufferQueueFunction(CreateBufferQueueFunction f) {
@@ -72,6 +66,9 @@
     auto& mutableEventThread() { return mFlinger->mEventThread; }
     auto& mutableEventQueue() { return mFlinger->mEventQueue; }
 
+    auto& mutableHwcDisplayData() { return mFlinger->getBE().mHwc->mDisplayData; }
+    auto& mutableHwcDisplaySlots() { return mFlinger->getBE().mHwc->mHwcDisplaySlots; }
+
     ~TestableSurfaceFlinger() {
         // All these pointer and container clears help ensure that GMock does
         // not report a leaked object, since the SurfaceFlinger instance may
@@ -83,6 +80,24 @@
         mFlinger->getBE().mRenderEngine.reset();
     }
 
+    /* ------------------------------------------------------------------------
+     * Wrapper classes for Read-write access to private data to set up
+     * preconditions and assert post-conditions.
+     */
+    struct HWC2Display : public HWC2::Display {
+        HWC2Display(Hwc2::Composer& composer,
+                    const std::unordered_set<HWC2::Capability>& capabilities, hwc2_display_t id,
+                    HWC2::DisplayType type)
+              : HWC2::Display(composer, capabilities, id, type) {}
+        ~HWC2Display() {
+            // Prevents a call to disable vsyncs.
+            mType = HWC2::DisplayType::Invalid;
+        }
+
+        auto& mutableIsConnected() { return this->mIsConnected; }
+        auto& mutableConfigs() { return this->mConfigs; }
+    };
+
     sp<SurfaceFlinger> mFlinger = new SurfaceFlinger();
 };