From 72762260877d4d7283179484b2565ccc8c1d4013 Mon Sep 17 00:00:00 2001
From: Chris Drury <cdrury@squareup.com>
Date: Mon, 26 Feb 2024 13:51:44 -0500
Subject: [PATCH] Use our own HandlerDispatcher for ComposeUi tests (#1306)

---
 gradle/libs.versions.toml                           |  4 +++-
 .../test/projects/coroutine-delay-main/build.gradle |  2 ++
 .../paparazzi/plugin/test/CoroutineDelayMainTest.kt |  9 +++++++++
 .../src/main/java/app/cash/paparazzi/Paparazzi.kt   | 13 +++++++++----
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 0c74ec653d..d30073cc36 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -3,6 +3,7 @@ agp = "8.2.2"
 androidTools = "31.2.2" # == 23.0.0 + agp version
 bytebuddy = "1.14.12"
 composeCompiler = "1.5.8"
+coroutines = "1.8.0"
 javaTarget = "11"
 jcodec = "0.2.5"
 kotlin = "1.9.22"
@@ -37,7 +38,8 @@ jcodec-core = { module = "org.jcodec:jcodec", version.ref = "jcodec" }
 jcodec-javase = { module = "org.jcodec:jcodec-javase", version.ref = "jcodec" }
 
 kotlin-bom = { module = "org.jetbrains.kotlin:kotlin-bom", version.ref = "kotlin" }
-kotlinx-coroutines-core = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-core", version = "1.8.0" }
+kotlinx-coroutines-core = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-core", version.ref = "coroutines" }
+kotlinx-coroutines-test = { module = "org.jetbrains.kotlinx:kotlinx-coroutines-test", version.ref = "coroutines" }
 
 ksp = { module = "com.google.devtools.ksp:symbol-processing-api", version.ref = "ksp" }
 
diff --git a/paparazzi-gradle-plugin/src/test/projects/coroutine-delay-main/build.gradle b/paparazzi-gradle-plugin/src/test/projects/coroutine-delay-main/build.gradle
index c2e4049839..8653bb4cac 100644
--- a/paparazzi-gradle-plugin/src/test/projects/coroutine-delay-main/build.gradle
+++ b/paparazzi-gradle-plugin/src/test/projects/coroutine-delay-main/build.gradle
@@ -27,6 +27,8 @@ android {
 
 dependencies {
   implementation libs.composeUi.material
+
+  testImplementation libs.kotlinx.coroutines.test
 }
 
 apply from: '../guava-fix.gradle'
diff --git a/paparazzi-gradle-plugin/src/test/projects/coroutine-delay-main/src/test/java/app/cash/paparazzi/plugin/test/CoroutineDelayMainTest.kt b/paparazzi-gradle-plugin/src/test/projects/coroutine-delay-main/src/test/java/app/cash/paparazzi/plugin/test/CoroutineDelayMainTest.kt
index c3a3433ac9..2d21a21435 100644
--- a/paparazzi-gradle-plugin/src/test/projects/coroutine-delay-main/src/test/java/app/cash/paparazzi/plugin/test/CoroutineDelayMainTest.kt
+++ b/paparazzi-gradle-plugin/src/test/projects/coroutine-delay-main/src/test/java/app/cash/paparazzi/plugin/test/CoroutineDelayMainTest.kt
@@ -7,12 +7,21 @@ import app.cash.paparazzi.Paparazzi
 import org.junit.Assert.assertEquals
 import org.junit.Rule
 import org.junit.Test
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.test.resetMain
 
+@OptIn(ExperimentalCoroutinesApi::class)
 class CoroutineDelayMainTest {
   @get:Rule
   val paparazzi = Paparazzi()
 
+  init {
+    // coroutines-test installs a TestMainDispatcher which is incompatible with ComposeUi
+    Dispatchers.resetMain()
+  }
+
   @Test fun delayUsesMainDispatcher() {
     var start = 0L
     var end = 0L
diff --git a/paparazzi/src/main/java/app/cash/paparazzi/Paparazzi.kt b/paparazzi/src/main/java/app/cash/paparazzi/Paparazzi.kt
index 655f87c3c6..5e2492bfbd 100644
--- a/paparazzi/src/main/java/app/cash/paparazzi/Paparazzi.kt
+++ b/paparazzi/src/main/java/app/cash/paparazzi/Paparazzi.kt
@@ -19,6 +19,7 @@ import android.animation.AnimationHandler
 import android.content.Context
 import android.content.res.Resources
 import android.graphics.Bitmap
+import android.os.Handler
 import android.os.Handler_Delegate
 import android.os.SystemClock_Delegate
 import android.util.AttributeSet
@@ -85,7 +86,7 @@ import java.awt.image.BufferedImage
 import java.util.Date
 import java.util.EnumSet
 import java.util.concurrent.TimeUnit
-import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.android.asCoroutineDispatcher
 
 @OptIn(ExperimentalComposeUiApi::class, InternalComposeUiApi::class)
 public class Paparazzi @JvmOverloads constructor(
@@ -303,9 +304,9 @@ public class Paparazzi @JvmOverloads constructor(
 
           // By default, Compose UI uses its own implementation of CoroutineDispatcher, `AndroidUiDispatcher`.
           // Since this dispatcher does not provide its own implementation of Delay, it will default to using DefaultDelay which runs
-          // async to our test Handler. By initializing Recomposer with Dispatchers.Main, Delay will now be backed by our test Handler,
-          // synchronizing expected behavior.
-          WindowRecomposerPolicy.setFactory { it.createLifecycleAwareWindowRecomposer(Dispatchers.Main) }
+          // async to our test Handler. By initializing Recomposer with our own HandlerDispatcher, Delay will now be backed
+          // by our test Handler, synchronizing expected behavior.
+          WindowRecomposerPolicy.setFactory { it.createLifecycleAwareWindowRecomposer(MAIN_DISPATCHER) }
         }
 
         if (hasLifecycleOwnerRuntime) {
@@ -647,6 +648,10 @@ public class Paparazzi @JvmOverloads constructor(
 
     internal lateinit var sessionParamsBuilder: SessionParamsBuilder
 
+    private val MAIN_DISPATCHER by lazy {
+      Handler.getMain().asCoroutineDispatcher("Paparazzi-Main")
+    }
+
     private val isVerifying: Boolean =
       System.getProperty("paparazzi.test.verify")?.toBoolean() == true