diff --git a/tests/Android.bp b/tests/Android.bp
index 4bc654c..fc08e86 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -168,6 +168,7 @@
         "src/**/*Test.java",
         "src/**/*Test.kt",
         "src/**/RoboApiWrapper.kt",
+        "src/**/EventsRule.kt",
         "multivalentTests/src/**/*Test.java",
         "multivalentTests/src/**/*Test.kt",
     ],
diff --git a/tests/src/com/android/launcher3/celllayout/integrationtest/events/EventWaiter.kt b/tests/src/com/android/launcher3/celllayout/integrationtest/events/EventWaiter.kt
new file mode 100644
index 0000000..20ad60f
--- /dev/null
+++ b/tests/src/com/android/launcher3/celllayout/integrationtest/events/EventWaiter.kt
@@ -0,0 +1,52 @@
+/*
+ * 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.celllayout.integrationtest.events
+
+import com.android.launcher3.debug.TestEventEmitter.TestEvent
+import java.util.concurrent.TimeUnit
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeoutOrNull
+
+enum class EventStatus() {
+    SUCCESS,
+    FAILURE,
+    TIMEOUT,
+}
+
+class EventWaiter(val eventToWait: TestEvent) {
+    private val deferrable = CompletableDeferred<EventStatus>()
+
+    companion object {
+        private const val TAG = "EventWaiter"
+        private val SIGNAL_TIMEOUT = TimeUnit.SECONDS.toMillis(5)
+    }
+
+    fun waitForSignal(timeout: Long = SIGNAL_TIMEOUT) = runBlocking {
+        var status = withTimeoutOrNull(timeout) { deferrable.await() }
+        if (status == null) {
+            status = EventStatus.TIMEOUT
+        }
+        if (status != EventStatus.SUCCESS) {
+            throw Exception("Failure waiting for event $eventToWait, failure = $status")
+        }
+    }
+
+    fun terminate() {
+        deferrable.complete(EventStatus.SUCCESS)
+    }
+}
diff --git a/tests/src/com/android/launcher3/celllayout/integrationtest/events/EventsRule.kt b/tests/src/com/android/launcher3/celllayout/integrationtest/events/EventsRule.kt
index fb61ced..45eb5e1 100644
--- a/tests/src/com/android/launcher3/celllayout/integrationtest/events/EventsRule.kt
+++ b/tests/src/com/android/launcher3/celllayout/integrationtest/events/EventsRule.kt
@@ -17,11 +17,15 @@
 package com.android.launcher3.celllayout.integrationtest.events
 
 import android.content.Context
-import com.android.launcher3.debug.TestEvent
+import android.util.Log
+import com.android.dx.mockito.inline.extended.ExtendedMockito.*
+import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import com.android.launcher3.debug.TestEventEmitter
+import com.android.launcher3.debug.TestEventEmitter.TestEvent
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
+import org.mockito.quality.Strictness
 
 /**
  * Rule to create EventWaiters to wait for events that happens on the Launcher. For reference look
@@ -30,35 +34,65 @@
  * Waiting for event should be used to prevent race conditions, it provides a more precise way of
  * waiting for events compared to [AbstractLauncherUiTest#waitForLauncherCondition].
  *
- * This class overrides the [TestEventEmitter] with [TestEventsEmitterImplementation] and makes sure
- * to return the [TestEventEmitter] to the previous value when finished.
+ * This class mocks the static method [TestEventEmitter.sendEvent]
  */
 class EventsRule(val context: Context) : TestRule {
 
-    private var prevEventEmitter: TestEventEmitter? = null
+    private val expectedEvents: ArrayDeque<EventWaiter> = ArrayDeque()
 
-    private val eventEmitter = TestEventsEmitterImplementation()
+    private lateinit var mockitoSession: StaticMockitoSession
 
     override fun apply(base: Statement, description: Description?): Statement {
         return object : Statement() {
             override fun evaluate() {
-                beforeTest()
-                base.evaluate()
-                afterTest()
+                try {
+                    beforeTest()
+                    base.evaluate()
+                } finally {
+                    afterTest()
+                }
             }
         }
     }
 
     fun createEventWaiter(expectedEvent: TestEvent): EventWaiter {
-        return eventEmitter.createEventWaiter(expectedEvent)
+        val eventWaiter = EventWaiter(expectedEvent)
+        expectedEvents.add(eventWaiter)
+        return eventWaiter
     }
 
     private fun beforeTest() {
-        prevEventEmitter = TestEventEmitter.INSTANCE.get(context)
-        TestEventEmitter.INSTANCE.initializeForTesting(eventEmitter)
+        mockitoSession =
+            mockitoSession()
+                .strictness(Strictness.LENIENT)
+                .spyStatic(TestEventEmitter::class.java)
+                .startMocking()
+
+        doAnswer { invocation ->
+                val event = (invocation.arguments[0] as TestEvent)
+                Log.d(TAG, "Signal received $event")
+                Log.d(TAG, "Total expected events ${expectedEvents.size}")
+                if (!expectedEvents.isEmpty()) {
+                    val eventWaiter = expectedEvents.last()
+                    if (eventWaiter.eventToWait == event) {
+                        Log.d(TAG, "Removing $event")
+                        expectedEvents.removeLast()
+                        eventWaiter.terminate()
+                    } else {
+                        Log.d(TAG, "Not matching $event")
+                    }
+                }
+                null
+            }
+            .`when` { TestEventEmitter.sendEvent(any()) }
     }
 
     private fun afterTest() {
-        TestEventEmitter.INSTANCE.initializeForTesting(prevEventEmitter)
+        mockitoSession.finishMocking()
+        expectedEvents.clear()
+    }
+
+    companion object {
+        private const val TAG = "TestEvents"
     }
 }
diff --git a/tests/src/com/android/launcher3/celllayout/integrationtest/events/TestEventsEmitterImplementation.kt b/tests/src/com/android/launcher3/celllayout/integrationtest/events/TestEventsEmitterImplementation.kt
deleted file mode 100644
index 5e062d0..0000000
--- a/tests/src/com/android/launcher3/celllayout/integrationtest/events/TestEventsEmitterImplementation.kt
+++ /dev/null
@@ -1,90 +0,0 @@
-/*
- * 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.celllayout.integrationtest.events
-
-import android.util.Log
-import com.android.launcher3.debug.TestEvent
-import com.android.launcher3.debug.TestEventEmitter
-import java.util.concurrent.TimeUnit
-import kotlinx.coroutines.CompletableDeferred
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.withTimeoutOrNull
-
-enum class EventStatus() {
-    SUCCESS,
-    FAILURE,
-    TIMEOUT,
-}
-
-class EventWaiter(val eventToWait: TestEvent) {
-    private val deferrable = CompletableDeferred<EventStatus>()
-
-    companion object {
-        private const val TAG = "EventWaiter"
-        private val SIGNAL_TIMEOUT = TimeUnit.SECONDS.toMillis(5)
-    }
-
-    fun waitForSignal(timeout: Long = SIGNAL_TIMEOUT) = runBlocking {
-        var status = withTimeoutOrNull(timeout) { deferrable.await() }
-        if (status == null) {
-            status = EventStatus.TIMEOUT
-        }
-        if (status != EventStatus.SUCCESS) {
-            throw Exception("Failure waiting for event $eventToWait, failure = $status")
-        }
-    }
-
-    fun terminate() {
-        deferrable.complete(EventStatus.SUCCESS)
-    }
-}
-
-class TestEventsEmitterImplementation() : TestEventEmitter {
-    companion object {
-        private const val TAG = "TestEvents"
-    }
-
-    private val expectedEvents: ArrayDeque<EventWaiter> = ArrayDeque()
-
-    fun createEventWaiter(expectedEvent: TestEvent): EventWaiter {
-        val eventWaiter = EventWaiter(expectedEvent)
-        expectedEvents.add(eventWaiter)
-        return eventWaiter
-    }
-
-    private fun clearQueue() {
-        expectedEvents.clear()
-    }
-
-    override fun sendEvent(event: TestEvent) {
-        Log.d(TAG, "Signal received $event")
-        Log.d(TAG, "Total expected events ${expectedEvents.size}")
-        if (expectedEvents.isEmpty()) return
-        val eventWaiter = expectedEvents.last()
-        if (eventWaiter.eventToWait == event) {
-            Log.d(TAG, "Removing $event")
-            expectedEvents.removeLast()
-            eventWaiter.terminate()
-        } else {
-            Log.d(TAG, "Not matching $event")
-        }
-    }
-
-    override fun close() {
-        clearQueue()
-    }
-}