From f369d51cb77ad1de3aad0f35bcea6ffb31734fab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=E2=89=A1ZRS?= <12814349+LZRS@users.noreply.github.com>
Date: Fri, 6 Dec 2024 13:06:09 +0300
Subject: [PATCH] Monitor questionnaire viewmodel dispatcher to track idling

---
 .../test/EspressoTrackedDispatcher.kt         | 49 +++++++++++++++++++
 .../test/QuestionnaireUiEspressoTest.kt       | 26 ++++++++++
 .../datacapture/QuestionnaireViewModel.kt     |  7 +--
 3 files changed, 76 insertions(+), 6 deletions(-)
 create mode 100644 datacapture/src/androidTest/java/com/google/android/fhir/datacapture/test/EspressoTrackedDispatcher.kt

diff --git a/datacapture/src/androidTest/java/com/google/android/fhir/datacapture/test/EspressoTrackedDispatcher.kt b/datacapture/src/androidTest/java/com/google/android/fhir/datacapture/test/EspressoTrackedDispatcher.kt
new file mode 100644
index 0000000000..fa9d0f0859
--- /dev/null
+++ b/datacapture/src/androidTest/java/com/google/android/fhir/datacapture/test/EspressoTrackedDispatcher.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * 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.google.android.fhir.datacapture.test
+
+import androidx.test.espresso.IdlingRegistry
+import androidx.test.espresso.idling.CountingIdlingResource
+import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.CoroutineDispatcher
+
+/** https://github.com/Kotlin/kotlinx.coroutines/issues/242#issuecomment-561503344 */
+class EspressoTrackedDispatcher(private val wrappedCoroutineDispatcher: CoroutineDispatcher) :
+  CoroutineDispatcher() {
+  private val counter: CountingIdlingResource =
+    CountingIdlingResource("EspressoTrackedDispatcher for $wrappedCoroutineDispatcher")
+
+  init {
+    IdlingRegistry.getInstance().register(counter)
+  }
+
+  override fun dispatch(context: CoroutineContext, block: Runnable) {
+    counter.increment()
+    val blockWithDecrement = Runnable {
+      try {
+        block.run()
+      } finally {
+        counter.decrement()
+      }
+    }
+    wrappedCoroutineDispatcher.dispatch(context, blockWithDecrement)
+  }
+
+  fun cleanUp() {
+    IdlingRegistry.getInstance().unregister(counter)
+  }
+}
diff --git a/datacapture/src/androidTest/java/com/google/android/fhir/datacapture/test/QuestionnaireUiEspressoTest.kt b/datacapture/src/androidTest/java/com/google/android/fhir/datacapture/test/QuestionnaireUiEspressoTest.kt
index 3ca8eb1b21..c103541b7a 100644
--- a/datacapture/src/androidTest/java/com/google/android/fhir/datacapture/test/QuestionnaireUiEspressoTest.kt
+++ b/datacapture/src/androidTest/java/com/google/android/fhir/datacapture/test/QuestionnaireUiEspressoTest.kt
@@ -42,6 +42,7 @@ import ca.uhn.fhir.context.FhirContext
 import ca.uhn.fhir.context.FhirVersionEnum
 import ca.uhn.fhir.parser.IParser
 import com.google.android.fhir.datacapture.QuestionnaireFragment
+import com.google.android.fhir.datacapture.questionnaireViewModelCoroutineContext
 import com.google.android.fhir.datacapture.test.utilities.clickIcon
 import com.google.android.fhir.datacapture.test.utilities.clickOnText
 import com.google.android.fhir.datacapture.validation.Invalid
@@ -58,6 +59,7 @@ import java.time.LocalDate
 import java.time.LocalDateTime
 import java.util.Calendar
 import java.util.Date
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.test.runTest
 import org.hamcrest.CoreMatchers
 import org.hl7.fhir.r4.model.DateTimeType
@@ -68,11 +70,19 @@ import org.junit.Assert
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
+import org.junit.rules.TestRule
+import org.junit.runner.Description
 import org.junit.runner.RunWith
+import org.junit.runners.model.Statement
 
 @RunWith(AndroidJUnit4::class)
 class QuestionnaireUiEspressoTest {
 
+  @Rule
+  @JvmField
+  var questionnaireViewModelCoroutineContextIdlerRule =
+    QuestionnaireViewModelCoroutineContextIdlerRule()
+
   @Rule
   @JvmField
   var activityScenarioRule: ActivityScenarioRule<TestActivity> =
@@ -745,3 +755,19 @@ class QuestionnaireUiEspressoTest {
     return testQuestionnaireFragment!!.getQuestionnaireResponse()
   }
 }
+
+class QuestionnaireViewModelCoroutineContextIdlerRule : TestRule {
+  override fun apply(base: Statement?, description: Description?): Statement =
+    object : Statement() {
+      override fun evaluate() {
+        val espressoTrackedDispatcherDefault = EspressoTrackedDispatcher(Dispatchers.Default)
+        questionnaireViewModelCoroutineContext = espressoTrackedDispatcherDefault
+        try {
+          base?.evaluate()
+        } finally {
+          espressoTrackedDispatcherDefault.cleanUp()
+          questionnaireViewModelCoroutineContext = Dispatchers.Default
+        }
+      }
+    }
+}
diff --git a/datacapture/src/main/java/com/google/android/fhir/datacapture/QuestionnaireViewModel.kt b/datacapture/src/main/java/com/google/android/fhir/datacapture/QuestionnaireViewModel.kt
index a2c92aec9e..4f97e23390 100644
--- a/datacapture/src/main/java/com/google/android/fhir/datacapture/QuestionnaireViewModel.kt
+++ b/datacapture/src/main/java/com/google/android/fhir/datacapture/QuestionnaireViewModel.kt
@@ -596,12 +596,8 @@ internal class QuestionnaireViewModel(application: Application, state: SavedStat
   }
 
   private val _questionnaireStateFlow: Flow<QuestionnaireState> =
-    combine(modificationCount, currentPageIndexFlow, isInReviewModeFlow) {
-        modCount,
-        pageIndex,
-        inReview,
+    combine(modificationCount, currentPageIndexFlow, isInReviewModeFlow) { _, _, _,
         ->
-        println(Triple(modCount, pageIndex, inReview))
         getQuestionnaireState()
       }
       .flowOn(questionnaireViewModelCoroutineContext)
@@ -612,7 +608,6 @@ internal class QuestionnaireViewModel(application: Application, state: SavedStat
       .withIndex()
       .onEach {
         if (it.index == 0) {
-          println("how many times called")
           initializeCalculatedExpressions()
           modificationCount.update { count -> count + 1 }
         }