diff --git a/services/surfaceflinger/tests/unittests/CompositionTest.cpp b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
index 60da70f..89c5d8a 100644
--- a/services/surfaceflinger/tests/unittests/CompositionTest.cpp
+++ b/services/surfaceflinger/tests/unittests/CompositionTest.cpp
@@ -793,6 +793,9 @@
 
         sp<L> layer = factory();
 
+        // Layer should be registered with scheduler.
+        EXPECT_EQ(1, test->mFlinger.scheduler()->layerHistorySize());
+
         Mock::VerifyAndClear(test->mComposer);
         Mock::VerifyAndClear(test->mRenderEngine);
         Mock::VerifyAndClear(test->mMessageQueue);
@@ -828,6 +831,10 @@
 
         test->mDisplay->getCompositionDisplay()->clearOutputLayers();
         test->mFlinger.mutableDrawingState().layersSortedByZ.clear();
+
+        // Layer should be unregistered with scheduler.
+        test->mFlinger.onMessageReceived(MessageQueue::INVALIDATE);
+        EXPECT_EQ(0, test->mFlinger.scheduler()->layerHistorySize());
     }
 };
 
diff --git a/services/surfaceflinger/tests/unittests/LayerHistoryTest.cpp b/services/surfaceflinger/tests/unittests/LayerHistoryTest.cpp
index 9a962bc..e93d31e 100644
--- a/services/surfaceflinger/tests/unittests/LayerHistoryTest.cpp
+++ b/services/surfaceflinger/tests/unittests/LayerHistoryTest.cpp
@@ -29,10 +29,10 @@
 
     LayerHistoryTest() { mFlinger.resetScheduler(mScheduler); }
 
-    LayerHistory& history() { return mScheduler->mutableLayerHistory(); }
-    const LayerHistory& history() const { return mScheduler->mutableLayerHistory(); }
+    LayerHistory& history() { return *mScheduler->mutableLayerHistory(); }
+    const LayerHistory& history() const { return *mScheduler->mutableLayerHistory(); }
 
-    size_t layerCount() const NO_THREAD_SAFETY_ANALYSIS { return history().mLayerInfos.size(); }
+    size_t layerCount() const { return mScheduler->layerHistorySize(); }
     size_t activeLayerCount() const NO_THREAD_SAFETY_ANALYSIS { return history().mActiveLayersEnd; }
 
     size_t frequentLayerCount(nsecs_t now) const NO_THREAD_SAFETY_ANALYSIS {
diff --git a/services/surfaceflinger/tests/unittests/TestableScheduler.h b/services/surfaceflinger/tests/unittests/TestableScheduler.h
index ae6aa89..40c00c4 100644
--- a/services/surfaceflinger/tests/unittests/TestableScheduler.h
+++ b/services/surfaceflinger/tests/unittests/TestableScheduler.h
@@ -21,6 +21,7 @@
 
 #include "Scheduler/DispSync.h"
 #include "Scheduler/EventThread.h"
+#include "Scheduler/LayerHistory.h"
 #include "Scheduler/Scheduler.h"
 
 namespace android {
@@ -28,18 +29,26 @@
 class TestableScheduler : public Scheduler {
 public:
     explicit TestableScheduler(const scheduler::RefreshRateConfigs& configs)
-          : Scheduler([](bool) {}, configs) {}
+          : Scheduler([](bool) {}, configs) {
+        mLayerHistory.emplace();
+    }
 
     TestableScheduler(std::unique_ptr<DispSync> primaryDispSync,
                       std::unique_ptr<EventControlThread> eventControlThread,
                       const scheduler::RefreshRateConfigs& configs)
-          : Scheduler(std::move(primaryDispSync), std::move(eventControlThread), configs) {}
+          : Scheduler(std::move(primaryDispSync), std::move(eventControlThread), configs) {
+        mLayerHistory.emplace();
+    }
 
     // Used to inject mock event thread.
     ConnectionHandle createConnection(std::unique_ptr<EventThread> eventThread) {
         return Scheduler::createConnection(std::move(eventThread));
     }
 
+    size_t layerHistorySize() const NO_THREAD_SAFETY_ANALYSIS {
+        return mLayerHistory->mLayerInfos.size();
+    }
+
     /* ------------------------------------------------------------------------
      * Read-write access to private data to set up preconditions and assert
      * post-conditions.