From e8211216a1dba4811f1bc08f38ad9d0ae701ff46 Mon Sep 17 00:00:00 2001
From: Dmitry Khalanskiy <Dmitry.Khalanskiy@jetbrains.com>
Date: Thu, 28 Oct 2021 14:54:57 +0300
Subject: [PATCH] Simplify TestCoroutineScope cleanup

---
 .../common/src/TestCoroutineScope.kt          | 24 +++++++++++--------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/kotlinx-coroutines-test/common/src/TestCoroutineScope.kt b/kotlinx-coroutines-test/common/src/TestCoroutineScope.kt
index a7de96b710..c9e578a7a2 100644
--- a/kotlinx-coroutines-test/common/src/TestCoroutineScope.kt
+++ b/kotlinx-coroutines-test/common/src/TestCoroutineScope.kt
@@ -45,19 +45,23 @@ private class TestCoroutineScopeImpl(
 
     override fun cleanupTestCoroutines() {
         val delayController = coroutineContext.delayController
-        if (delayController != null) {
-            delayController.cleanupTestCoroutines()
-            coroutineContext.uncaughtExceptionCaptor.cleanupTestCoroutinesCaptor()
+        val hasUnfinishedJobs = if (delayController != null) {
+            try {
+                delayController.cleanupTestCoroutines()
+                false
+            } catch (e: UncompletedCoroutinesError) {
+                true
+            }
         } else {
             testScheduler.runCurrent()
-            coroutineContext.uncaughtExceptionCaptor.cleanupTestCoroutinesCaptor()
-            if (!testScheduler.isIdle()) {
-                throw UncompletedCoroutinesError(
-                    "Unfinished coroutines during teardown. Ensure all coroutines are" +
-                        " completed or cancelled by your test."
-                )
-            }
+            !testScheduler.isIdle()
         }
+        coroutineContext.uncaughtExceptionCaptor.cleanupTestCoroutinesCaptor()
+        if (hasUnfinishedJobs)
+            throw UncompletedCoroutinesError(
+                "Unfinished coroutines during teardown. Ensure all coroutines are" +
+                    " completed or cancelled by your test."
+            )
         val jobs = coroutineContext.activeJobs()
         if ((jobs - initialJobs).isNotEmpty())
             throw UncompletedCoroutinesError("Test finished with active jobs: $jobs")