diff --git a/src/com/android/launcher3/util/ScreenOnTracker.java b/src/com/android/launcher3/util/ScreenOnTracker.java
index 12eff61..8ee799a 100644
--- a/src/com/android/launcher3/util/ScreenOnTracker.java
+++ b/src/com/android/launcher3/util/ScreenOnTracker.java
@@ -24,7 +24,10 @@
 import android.content.Context;
 import android.content.Intent;
 
+import androidx.annotation.VisibleForTesting;
+
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.function.Consumer;
 
 /**
  * Utility class for tracking if the screen is currently on or off
@@ -34,8 +37,7 @@
     public static final MainThreadInitializedObject<ScreenOnTracker> INSTANCE =
             new MainThreadInitializedObject<>(ScreenOnTracker::new);
 
-    private final SimpleBroadcastReceiver mReceiver =
-            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::onReceive);
+    private final SimpleBroadcastReceiver mReceiver;
     private final CopyOnWriteArrayList<ScreenOnListener> mListeners = new CopyOnWriteArrayList<>();
 
     private final Context mContext;
@@ -44,8 +46,20 @@
     private ScreenOnTracker(Context context) {
         // Assume that the screen is on to begin with
         mContext = context;
+        mReceiver = new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::onReceive);
+        init();
+    }
+
+    @VisibleForTesting
+    ScreenOnTracker(Context context, SimpleBroadcastReceiver receiver) {
+        mContext = context;
+        mReceiver = receiver;
+        init();
+    }
+
+    private void init() {
         mIsScreenOn = true;
-        mReceiver.register(context, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
+        mReceiver.register(mContext, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
     }
 
     @Override
@@ -53,7 +67,8 @@
         mReceiver.unregisterReceiverSafely(mContext);
     }
 
-    private void onReceive(Intent intent) {
+    @VisibleForTesting
+    void onReceive(Intent intent) {
         String action = intent.getAction();
         if (ACTION_SCREEN_ON.equals(action)) {
             mIsScreenOn = true;
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/ScreenOnTrackerTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/ScreenOnTrackerTest.kt
new file mode 100644
index 0000000..430aad2
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/ScreenOnTrackerTest.kt
@@ -0,0 +1,103 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+package com.android.launcher3.util
+
+import android.content.Context
+import android.content.Intent
+import android.content.Intent.ACTION_SCREEN_OFF
+import android.content.Intent.ACTION_SCREEN_ON
+import android.content.Intent.ACTION_USER_PRESENT
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.verifyNoMoreInteractions
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ScreenOnTrackerTest {
+
+    @Mock private lateinit var receiver: SimpleBroadcastReceiver
+    @Mock private lateinit var context: Context
+    @Mock private lateinit var listener: ScreenOnTracker.ScreenOnListener
+
+    private lateinit var underTest: ScreenOnTracker
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        underTest = ScreenOnTracker(context, receiver)
+    }
+
+    @Test
+    fun test_default_state() {
+        verify(receiver).register(context, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT)
+        assertThat(underTest.isScreenOn).isTrue()
+    }
+
+    @Test
+    fun close_unregister_receiver() {
+        underTest.close()
+
+        verify(receiver).unregisterReceiverSafely(context)
+    }
+
+    @Test
+    fun add_listener_then_receive_screen_on_intent_notify_listener() {
+        underTest.addListener(listener)
+
+        underTest.onReceive(Intent(ACTION_SCREEN_ON))
+
+        verify(listener).onScreenOnChanged(true)
+        assertThat(underTest.isScreenOn).isTrue()
+    }
+
+    @Test
+    fun add_listener_then_receive_screen_off_intent_notify_listener() {
+        underTest.addListener(listener)
+
+        underTest.onReceive(Intent(ACTION_SCREEN_OFF))
+
+        verify(listener).onScreenOnChanged(false)
+        assertThat(underTest.isScreenOn).isFalse()
+    }
+
+    @Test
+    fun add_listener_then_receive_user_present_intent_notify_listener() {
+        underTest.addListener(listener)
+
+        underTest.onReceive(Intent(ACTION_USER_PRESENT))
+
+        verify(listener).onUserPresent()
+        assertThat(underTest.isScreenOn).isTrue()
+    }
+
+    @Test
+    fun remove_listener_then_receive_intent_noOp() {
+        underTest.addListener(listener)
+
+        underTest.removeListener(listener)
+
+        underTest.onReceive(Intent(ACTION_USER_PRESENT))
+        verifyNoMoreInteractions(listener)
+    }
+}