Improve test robustness (#4369)

* Fix testFutureIsDoneAfterChildrenCompleted flakiness

`invokeOnCompletion` can run in parallel to the code that has
already detected a coroutine's completion.
This has happened on the CI once.
The assumption that this can't happen has been removed from the
test.

* Improve println() detection in tests

Now, the printed messages are also preserved, and also, failing to
remove a println capturer from a test no longer affects the
unrelated tests.

* Move robust test cleanup

Author: dkhalanskyjb

Diff for: integration/kotlinx-coroutines-guava/test/ListenableFutureTest.kt

@@ -627,26 +627,24 @@ class ListenableFutureTest : TestBase() {
     fun testFutureIsDoneAfterChildrenCompleted() = runTest {
         expect(1)
         val testException = TestException()
-        val latch = CountDownLatch(1)
+        val futureIsAllowedToFinish = CountDownLatch(1)
         // Don't propagate exception to the test and use different dispatchers as we are going to block test thread.
         val future = future(context = NonCancellable + Dispatchers.Default) {
             val foo = async(start = CoroutineStart.UNDISPATCHED) {
                 try {
                     delay(Long.MAX_VALUE)
                     42
                 } finally {
-                    latch.await()
+                    futureIsAllowedToFinish.await()
+                    expect(3)
                 }
             }
-            foo.invokeOnCompletion {
-                expect(3)
-            }
             val bar = async<Int> { throw testException }
             foo.await() + bar.await()
         }
         yield()
         expect(2)
-        latch.countDown()
+        futureIsAllowedToFinish.countDown()
         // Blocking get should succeed after internal coroutine completes.
         val thrown = assertFailsWith<ExecutionException> { future.get() }
         expect(4)

Diff for: kotlinx-coroutines-debug/test/DebugTestBase.kt

@@ -14,18 +14,13 @@ open class DebugTestBase : TestBase() {
 
     @Before
     open fun setUp() {
-        before()
         DebugProbes.sanitizeStackTraces = false
         DebugProbes.enableCreationStackTraces = false
         DebugProbes.install()
     }
 
     @After
     fun tearDown() {
-        try {
-            DebugProbes.uninstall()
-        } finally {
-            onCompletion()
-        }
+        DebugProbes.uninstall()
     }
 }

Diff for: kotlinx-coroutines-debug/test/ToStringTest.kt

@@ -9,24 +9,7 @@ import java.io.*
 import kotlin.coroutines.*
 import kotlin.test.*
 
-class ToStringTest : TestBase() {
-
-    @Before
-    fun setUp() {
-        before()
-        DebugProbes.sanitizeStackTraces = false
-        DebugProbes.install()
-    }
-
-    @After
-    fun tearDown() {
-        try {
-            DebugProbes.uninstall()
-        } finally {
-            onCompletion()
-        }
-    }
-
+class ToStringTest : DebugTestBase() {
 
     private suspend fun CoroutineScope.launchNestedScopes(): Job {
         return launch {

Diff for: test-utils/common/src/TestBase.common.kt

@@ -133,7 +133,9 @@ interface ErrorCatching {
         fun close() {
             synchronized(lock) {
                 if (closed) {
-                    lastResortReportException(IllegalStateException("ErrorCatching closed more than once"))
+                    val error = IllegalStateException("ErrorCatching closed more than once")
+                    lastResortReportException(error)
+                    errors.add(error)
                 }
                 closed = true
                 errors.firstOrNull()?.let {

Diff for: test-utils/jvm/src/TestBase.kt

@@ -5,6 +5,7 @@ import java.io.*
 import java.util.*
 import kotlin.coroutines.*
 import kotlinx.coroutines.*
+import java.util.concurrent.atomic.AtomicReference
 import kotlin.test.*
 
 actual val VERBOSE = try {
@@ -68,23 +69,9 @@ actual open class TestBase(
     private lateinit var threadsBefore: Set<Thread>
     private val uncaughtExceptions = Collections.synchronizedList(ArrayList<Throwable>())
     private var originalUncaughtExceptionHandler: Thread.UncaughtExceptionHandler? = null
-    /*
-     * System.out that we redefine in order to catch any debugging/diagnostics
-     * 'println' from main source set.
-     * NB: We do rely on the name 'previousOut' in the FieldWalker in order to skip its
-     * processing
-     */
-    private lateinit var previousOut: PrintStream
-
-    private object TestOutputStream : PrintStream(object : OutputStream() {
-        override fun write(b: Int) {
-            error("Detected unexpected call to 'println' from source code")
-        }
-    })
 
     actual fun println(message: Any?) {
-        if (disableOutCheck) kotlin.io.println(message)
-        else previousOut.println(message)
+        PrintlnStrategy.actualSystemOut.println(message)
     }
 
     @BeforeTest
@@ -97,34 +84,33 @@ actual open class TestBase(
             e.printStackTrace()
             uncaughtExceptions.add(e)
         }
-        if (!disableOutCheck) {
-            previousOut = System.out
-            System.setOut(TestOutputStream)
-        }
+        PrintlnStrategy.configure(disableOutCheck)
     }
 
     @AfterTest
     fun onCompletion() {
         // onCompletion should not throw exceptions before it finishes all cleanup, so that other tests always
-        // start in a clear, restored state
-        checkFinishCall()
-        if (!disableOutCheck) { // Restore global System.out first
-            System.setOut(previousOut)
+        // start in a clear, restored state, so we postpone throwing the observed errors.
+        fun cleanupStep(block: () -> Unit) {
+            try {
+                block()
+            } catch (e: Throwable) {
+                reportError(e)
+            }
         }
+        cleanupStep { checkFinishCall() }
+        // Reset the output stream first
+        cleanupStep { PrintlnStrategy.reset() }
         // Shutdown all thread pools
-        shutdownPoolsAfterTest()
+        cleanupStep { shutdownPoolsAfterTest() }
         // Check that are now leftover threads
-        runCatching {
-            checkTestThreads(threadsBefore)
-        }.onFailure {
-            reportError(it)
-        }
+        cleanupStep { checkTestThreads(threadsBefore) }
         // Restore original uncaught exception handler after the main shutdown sequence
         Thread.setDefaultUncaughtExceptionHandler(originalUncaughtExceptionHandler)
         if (uncaughtExceptions.isNotEmpty()) {
-            error("Expected no uncaught exceptions, but got $uncaughtExceptions")
+            reportError(IllegalStateException("Expected no uncaught exceptions, but got $uncaughtExceptions"))
         }
-        // The very last action -- throw error if any was detected
+        // The very last action -- throw all the detected errors
         errorCatching.close()
     }
 
@@ -164,6 +150,81 @@ actual open class TestBase(
     protected suspend fun currentDispatcher() = coroutineContext[ContinuationInterceptor]!!
 }
 
+private object PrintlnStrategy {
+    /**
+     * Installs a custom [PrintStream] instead of [System.out] to capture all the output and throw an exception if
+     * any was detected.
+     *
+     * Removes the previously set println handler and throws the exceptions detected by it.
+     * If [disableOutCheck] is set, this is the only effect.
+     */
+    fun configure(disableOutCheck: Boolean) {
+        val systemOut = System.out
+        if (systemOut is TestOutputStream) {
+            try {
+                systemOut.remove()
+            } catch (e: AssertionError) {
+                throw AssertionError("The previous TestOutputStream contained ", e)
+            }
+        }
+        if (!disableOutCheck) {
+            // Invariant: at most one indirection level in `TestOutputStream`.
+            System.setOut(TestOutputStream(actualSystemOut))
+        }
+    }
+
+    /**
+     * Removes the custom [PrintStream] and throws an exception if any output was detected.
+     */
+    fun reset() {
+        (System.out as? TestOutputStream)?.remove()
+    }
+
+    /**
+     * The [PrintStream] representing the actual stdout, ignoring the replacement [TestOutputStream].
+     */
+    val actualSystemOut: PrintStream get() = when (val out = System.out) {
+        is TestOutputStream -> out.previousOut
+        else -> out
+    }
+
+    private class TestOutputStream(
+        /*
+         * System.out that we redefine in order to catch any debugging/diagnostics
+         * 'println' from main source set.
+         * NB: We do rely on the name 'previousOut' in the FieldWalker in order to skip its
+         * processing
+         */
+        val previousOut: PrintStream,
+        private val myOutputStream: MyOutputStream = MyOutputStream(),
+    ) : PrintStream(myOutputStream) {
+
+        fun remove() {
+            System.setOut(previousOut)
+            if (myOutputStream.firstPrintStacktace.get() != null) {
+                throw AssertionError(
+                    "Detected a println. The captured output is: <<<${myOutputStream.capturedOutput}>>>",
+                    myOutputStream.firstPrintStacktace.get()
+                )
+            }
+        }
+
+        private class MyOutputStream(): OutputStream() {
+            val capturedOutput = ByteArrayOutputStream()
+
+            val firstPrintStacktace = AtomicReference<Throwable?>(null)
+
+            override fun write(b: Int) {
+                if (firstPrintStacktace.get() == null) {
+                    firstPrintStacktace.compareAndSet(null, IllegalStateException())
+                }
+                capturedOutput.write(b)
+            }
+        }
+
+    }
+}
+
 @Suppress("INVISIBLE_MEMBER", "INVISIBLE_REFERENCE")
 fun initPoolsBeforeTest() {
     DefaultScheduler.usePrivateScheduler()