diff --git a/sample/src/androidTest/java/com/airbnb/lottie/samples/FragmentVisibilityTests.kt b/sample/src/androidTest/java/com/airbnb/lottie/samples/FragmentVisibilityTests.kt
index 2bed7c9..f8f8001 100644
--- a/sample/src/androidTest/java/com/airbnb/lottie/samples/FragmentVisibilityTests.kt
+++ b/sample/src/androidTest/java/com/airbnb/lottie/samples/FragmentVisibilityTests.kt
@@ -119,7 +119,7 @@
             val animationView = fragment.requireView().findViewById<LottieAnimationView>(R.id.animation_view)
             // Wait for the animation view.
             // We have to use a property reference because the Fragment isn't resumed.
-            assertTrue(fragment.animationView.isAnimating)
+            assertTrue(animationView.isAnimating)
         }
     }
 
@@ -331,8 +331,6 @@
                                     animationWasPlayed = true
                                     IdlingRegistry.getInstance().register(LottieIdlingResource(this, name = "Lottie ${Random.nextFloat()}"))
                                 }
-                            } else {
-                                IdlingRegistry.getInstance().register(LottieIdlingAnimationResource(animationView, name = "Lottie finished animation ${Random.nextFloat()}"))
                             }
                         }
 
@@ -352,7 +350,10 @@
         scenario.onFragment { it.requireView().scrollBy(0, -10_000) }
         scenario.onFragment { assertTrue(it.animationView!!.isAnimating) }
         onView(withId(R.id.animation_view)).check(matches(isDisplayed()))
-        scenario.onFragment { assertFalse(it.animationView!!.isAnimating) }
+        scenario.onAnimationEnded()
+        scenario.onFragment { fragment ->
+            assertFalse(fragment.animationView!!.isAnimating)
+        }
     }
 
     @Test
@@ -395,8 +396,6 @@
                                     animationWasPlayed = true
                                     IdlingRegistry.getInstance().register(LottieIdlingResource(this, name = "Lottie ${Random.nextFloat()}"))
                                 }
-                            } else {
-                                IdlingRegistry.getInstance().register(LottieIdlingAnimationResource(animationView, name = "Lottie finished animation ${Random.nextFloat()}"))
                             }
                         }
 
@@ -415,7 +414,7 @@
         scenario.onFragment { assertFalse(it.animationView!!.isAnimating) }
         scenario.onFragment { it.requireView().scrollBy(0, -10_000) }
         scenario.onFragment { assertTrue(it.animationView!!.isAnimating) }
-        onView(withId(R.id.animation_view)).check(matches(isDisplayed()))
+        scenario.onAnimationEnded()
         scenario.onFragment { assertFalse(it.animationView!!.isAnimating) }
         scenario.onFragment { it.requireView().scrollBy(0, 10_000) }
         scenario.onFragment { it.requireView().scrollBy(0, -10_000) }
@@ -460,5 +459,12 @@
         }
     }
 
+    private fun <T : Fragment> FragmentScenario<T>.onAnimationEnded() {
+        onFragment { fragment ->
+            IdlingRegistry.getInstance().register(LottieIdlingAnimationResource(fragment.animationView!!, name = "Lottie finished animation ${Random.nextFloat()}"))
+        }
+        onIdle()
+    }
+
     private val Fragment.animationView get() = requireView().findViewById<LottieAnimationView>(R.id.animation_view)
 }
\ No newline at end of file
diff --git a/sample/src/androidTest/java/com/airbnb/lottie/samples/LottieIdlingAnimationResource.kt b/sample/src/androidTest/java/com/airbnb/lottie/samples/LottieIdlingAnimationResource.kt
index 1cac096..073ee80 100644
--- a/sample/src/androidTest/java/com/airbnb/lottie/samples/LottieIdlingAnimationResource.kt
+++ b/sample/src/androidTest/java/com/airbnb/lottie/samples/LottieIdlingAnimationResource.kt
@@ -5,16 +5,15 @@
 import androidx.test.espresso.IdlingResource
 import com.airbnb.lottie.LottieAnimationView
 
-class LottieIdlingAnimationResource(animationView: LottieAnimationView?, private val name: String = "Lottie") : IdlingResource {
+class LottieIdlingAnimationResource(animationView: LottieAnimationView, private val name: String = "Lottie") : IdlingResource {
+
+    private var hasEnded = false
+    private var callback: IdlingResource.ResourceCallback? = null
 
     init {
-        animationView?.addAnimatorListener(object : AnimatorListenerAdapter() {
-            override fun onAnimationStart(animation: Animator) {
-                isIdle = false
-            }
-
+        animationView.addAnimatorListener(object : AnimatorListenerAdapter() {
             override fun onAnimationEnd(animation: Animator) {
-                isIdle = true
+                hasEnded = true
                 callback?.onTransitionToIdle()
                 animationView.removeAllAnimatorListeners()
                 IdlingRegistry.getInstance().unregister(this@LottieIdlingAnimationResource)
@@ -22,16 +21,13 @@
         })
     }
 
-    private var callback: IdlingResource.ResourceCallback? = null
-    private var isIdle = animationView?.isAnimating?.not() ?: true
-
 
     override fun getName() = name
 
-    override fun isIdleNow() = isIdle
+    override fun isIdleNow() = hasEnded
 
     override fun registerIdleTransitionCallback(callback: IdlingResource.ResourceCallback?) {
         this.callback = callback
-        if (isIdle) callback?.onTransitionToIdle()
+        if (isIdleNow) callback?.onTransitionToIdle()
     }
 }
\ No newline at end of file