diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromFixedOrientationTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromFixedOrientationTest.kt
index 44bd8c8..52e9652 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromFixedOrientationTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/ShowImeOnAppStartWhenLaunchingAppFromFixedOrientationTest.kt
@@ -19,9 +19,7 @@
 import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.tools.common.Rotation
-import android.tools.common.Timestamp
-import android.tools.common.flicker.subject.exceptions.ExceptionMessageBuilder
-import android.tools.common.flicker.subject.exceptions.InvalidPropertyException
+import android.tools.common.flicker.subject.layers.LayerTraceEntrySubject
 import android.tools.common.traces.component.ComponentNameMatcher
 import android.tools.device.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.device.flicker.legacy.FlickerBuilder
@@ -85,44 +83,37 @@
     @Postsubmit
     @Test
     fun imeLayerAlphaOneAfterSnapshotStartingWindowRemoval() {
-        // Check if the snapshot appeared during the trace
-        var imeSnapshotRemovedTimestamp: Timestamp? = null
+        val layerTrace = flicker.reader.readLayersTrace() ?: error("Unable to read layers trace")
 
-        val layerTrace = flicker.reader.readLayersTrace()
-        val layerTraceEntries = layerTrace?.entries?.toList() ?: emptyList()
-
-        layerTraceEntries.zipWithNext { prev, next ->
-            val prevSnapshotLayerVisible =
-                ComponentNameMatcher.SNAPSHOT.layerMatchesAnyOf(prev.visibleLayers)
-            val nextSnapshotLayerVisible =
-                ComponentNameMatcher.SNAPSHOT.layerMatchesAnyOf(next.visibleLayers)
-
-            if (
-                imeSnapshotRemovedTimestamp == null &&
-                    (prevSnapshotLayerVisible && !nextSnapshotLayerVisible)
-            ) {
-                imeSnapshotRemovedTimestamp = next.timestamp
+        // Find the entries immediately after the IME snapshot has disappeared
+        val imeSnapshotRemovedEntries = layerTrace.entries.asSequence()
+            .zipWithNext { prev, next ->
+                if (ComponentNameMatcher.SNAPSHOT.layerMatchesAnyOf(prev.visibleLayers) &&
+                    !ComponentNameMatcher.SNAPSHOT.layerMatchesAnyOf(next.visibleLayers)) {
+                    next
+                } else {
+                    null
+                }
             }
-        }
+            .filterNotNull()
 
-        // if so, make an assertion
-        imeSnapshotRemovedTimestamp?.let { timestamp ->
-            val stateAfterSnapshot =
-                layerTrace?.getEntryAt(timestamp) ?: error("State not found for $timestamp")
+        // If we find it, make sure the IME is visible and fully animated in.
+        imeSnapshotRemovedEntries.forEach { entry ->
+            val entrySubject = LayerTraceEntrySubject(entry)
+            val imeLayerSubjects = entrySubject.subjects.filter {
+                ComponentNameMatcher.IME.layerMatchesAnyOf(it.layer) && it.isVisible
+            }
 
-            val imeLayers =
-                ComponentNameMatcher.IME.filterLayers(stateAfterSnapshot.visibleLayers.toList())
+            entrySubject
+                .check { "InputMethod must exist and be visible" }
+                .that(imeLayerSubjects.isNotEmpty())
+                .isEqual(true)
 
-            require(imeLayers.isNotEmpty()) { "IME layer not found" }
-            if (imeLayers.any { it.color.a != 1.0f }) {
-                val errorMsgBuilder =
-                    ExceptionMessageBuilder()
-                        .setTimestamp(timestamp)
-                        .forInvalidProperty("IME layer alpha")
-                        .setExpected("is 1.0")
-                        .setActual("not 1.0")
-                        .addExtraDescription("Filter", ComponentNameMatcher.IME.toLayerIdentifier())
-                throw InvalidPropertyException(errorMsgBuilder)
+            imeLayerSubjects.forEach { imeLayerSubject ->
+                imeLayerSubject
+                    .check { "alpha" }
+                    .that(imeLayerSubject.layer.color.a)
+                    .isEqual(1.0f)
             }
         }
     }