diff --git a/services/surfaceflinger/tests/fakehwc/Android.bp b/services/surfaceflinger/tests/fakehwc/Android.bp
index 68cf61e..a2c0611 100644
--- a/services/surfaceflinger/tests/fakehwc/Android.bp
+++ b/services/surfaceflinger/tests/fakehwc/Android.bp
@@ -31,7 +31,6 @@
         "libutils",
     ],
     static_libs: [
-        "libcompositionengine",
         "libgmock",
         "librenderengine",
         "libtrace_proto",
diff --git a/services/surfaceflinger/tests/unittests/CompositionTest.cpp b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
index 3fb8708..3addd61 100644
--- a/services/surfaceflinger/tests/unittests/CompositionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
@@ -17,6 +17,7 @@
 #undef LOG_TAG
 #define LOG_TAG "CompositionTest"
 
+#include <compositionengine/Display.h>
 #include <compositionengine/mock/DisplaySurface.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -753,12 +754,16 @@
     }
 
     static void injectLayer(CompositionTest* test, sp<Layer> layer) {
+        std::vector<std::unique_ptr<compositionengine::OutputLayer>> outputLayers;
+        outputLayers.emplace_back(
+                test->mDisplay->getCompositionDisplay()
+                        ->getOrCreateOutputLayer(layer->getCompositionLayer(), layer));
+        test->mDisplay->getCompositionDisplay()->setOutputLayersOrderedByZ(std::move(outputLayers));
+
         EXPECT_CALL(*test->mComposer, createLayer(HWC_DISPLAY, _))
                 .WillOnce(DoAll(SetArgPointee<1>(HWC_LAYER), Return(Error::NONE)));
 
-        const auto displayId = test->mDisplay->getId();
-        ASSERT_TRUE(displayId);
-        layer->createHwcLayer(&test->mFlinger.getHwComposer(), *displayId);
+        layer->createHwcLayer(&test->mFlinger.getHwComposer(), test->mDisplay);
 
         Mock::VerifyAndClear(test->mComposer);
 
@@ -771,10 +776,12 @@
     static void cleanupInjectedLayers(CompositionTest* test) {
         EXPECT_CALL(*test->mComposer, destroyLayer(HWC_DISPLAY, HWC_LAYER))
                 .WillOnce(Return(Error::NONE));
-        const auto displayId = test->mDisplay->getId();
-        ASSERT_TRUE(displayId);
+
+        test->mDisplay->getCompositionDisplay()->setOutputLayersOrderedByZ(
+                std::vector<std::unique_ptr<compositionengine::OutputLayer>>());
+
         for (auto layer : test->mFlinger.mutableDrawingState().layersSortedByZ) {
-            layer->destroyHwcLayer(*displayId);
+            layer->destroyHwcLayer(test->mDisplay);
         }
         test->mFlinger.mutableDrawingState().layersSortedByZ.clear();
     }
@@ -963,8 +970,8 @@
 };
 
 struct ForcedClientCompositionResultVariant : public RECompositionResultVariant {
-    static void setupLayerState(CompositionTest*, sp<Layer> layer) {
-        layer->forceClientComposition(DEFAULT_DISPLAY_ID);
+    static void setupLayerState(CompositionTest* test, sp<Layer> layer) {
+        layer->forceClientComposition(test->mDisplay);
     }
 
     template <typename Case>
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index e639b4d..bfdd20b 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -16,7 +16,10 @@
 
 #pragma once
 
+#include <compositionengine/Display.h>
+#include <compositionengine/OutputLayer.h>
 #include <compositionengine/impl/CompositionEngine.h>
+#include <compositionengine/impl/OutputLayerCompositionState.h>
 
 #include "BufferQueueLayer.h"
 #include "BufferStateLayer.h"
@@ -206,6 +209,14 @@
         layer->getBE().compositionInfo.hwc.sidebandStream = sidebandStream;
     }
 
+    void setLayerCompositionType(sp<Layer> layer, HWC2::Composition type) {
+        auto outputLayer = layer->findOutputLayerForDisplay(mFlinger->getDefaultDisplayDevice());
+        LOG_ALWAYS_FATAL_IF(!outputLayer);
+        auto& state = outputLayer->editState();
+        LOG_ALWAYS_FATAL_IF(!outputLayer->getState().hwc);
+        (*state.hwc).hwcCompositionType = static_cast<Hwc2::IComposerClient::Composition>(type);
+    };
+
     void setLayerPotentialCursor(sp<Layer> layer, bool potentialCursor) {
         layer->mPotentialCursor = potentialCursor;
     }