Generalize exceptions thrown from 'runWithTimeoutDumpingCoroutines'

Author: dkhalanskyjb

kotlinx-coroutines-debug/src/junit/CoroutinesTimeoutImpl.kt

@@ -4,7 +4,6 @@
 
 package kotlinx.coroutines.debug
 
-import org.junit.runners.model.*
 import java.util.concurrent.*
 
 /**
@@ -17,6 +16,7 @@ internal inline fun <T : Any?> runWithTimeoutDumpingCoroutines(
     methodName: String,
     testTimeoutMs: Long,
     cancelOnTimeout: Boolean,
+    initCancellationException: () -> Throwable,
     crossinline invocation: () -> T
 ): T {
     val testStartedLatch = CountDownLatch(1)
@@ -35,13 +35,14 @@ internal inline fun <T : Any?> runWithTimeoutDumpingCoroutines(
         testStartedLatch.await()
         return testResult.get(testTimeoutMs, TimeUnit.MILLISECONDS)
     } catch (e: TimeoutException) {
-        handleTimeout(testThread, methodName, testTimeoutMs, cancelOnTimeout)
+        handleTimeout(testThread, methodName, testTimeoutMs, cancelOnTimeout, initCancellationException())
     } catch (e: ExecutionException) {
         throw e.cause ?: e
     }
 }
 
-private fun handleTimeout(testThread: Thread, methodName: String, testTimeoutMs: Long, cancelOnTimeout: Boolean): Nothing {
+private fun handleTimeout(testThread: Thread, methodName: String, testTimeoutMs: Long, cancelOnTimeout: Boolean,
+                          cancellationException: Throwable): Nothing {
     val units =
         if (testTimeoutMs % 1000 == 0L)
             "${testTimeoutMs / 1000} seconds"
@@ -59,10 +60,11 @@ private fun handleTimeout(testThread: Thread, methodName: String, testTimeoutMs:
      * 2) Cancel all coroutines via debug agent API (changing system state!)
      * 3) Throw created exception
      */
-    val exception = createTimeoutException(testThread, testTimeoutMs)
+    cancellationException.attachStacktraceFrom(testThread)
+    testThread.interrupt()
     cancelIfNecessary(cancelOnTimeout)
     // If timed out test throws an exception, we can't do much except ignoring it
-    throw exception
+    throw cancellationException
 }
 
 private fun cancelIfNecessary(cancelOnTimeout: Boolean) {
@@ -73,10 +75,7 @@ private fun cancelIfNecessary(cancelOnTimeout: Boolean) {
     }
 }
 
-private fun createTimeoutException(thread: Thread, testTimeoutMs: Long): Exception {
+private fun Throwable.attachStacktraceFrom(thread: Thread) {
     val stackTrace = thread.stackTrace
-    val exception = TestTimedOutException(testTimeoutMs, TimeUnit.MILLISECONDS)
-    exception.stackTrace = stackTrace
-    thread.interrupt()
-    return exception
+    this.stackTrace = stackTrace
 }

kotlinx-coroutines-debug/src/junit/junit4/CoroutinesTimeoutStatement.kt

@@ -7,6 +7,7 @@ package kotlinx.coroutines.debug.junit4
 import kotlinx.coroutines.debug.*
 import org.junit.runner.*
 import org.junit.runners.model.*
+import java.util.concurrent.*
 
 internal class CoroutinesTimeoutStatement(
     private val testStatement: Statement,
@@ -17,7 +18,9 @@ internal class CoroutinesTimeoutStatement(
 
     override fun evaluate() {
         try {
-            runWithTimeoutDumpingCoroutines(testDescription.methodName, testTimeoutMs, cancelOnTimeout) {
+            runWithTimeoutDumpingCoroutines(testDescription.methodName, testTimeoutMs, cancelOnTimeout,
+                { TestTimedOutException(testTimeoutMs, TimeUnit.MILLISECONDS) })
+            {
                 testStatement.evaluate()
             }
         } finally {

kotlinx-coroutines-debug/src/junit/junit5/CoroutinesTimeoutExtension.kt

@@ -10,13 +10,9 @@ import org.junit.jupiter.api.extension.ExtensionContext
 import org.junit.jupiter.api.extension.InvocationInterceptor
 import org.junit.jupiter.api.extension.ReflectiveInvocationContext
 import org.junit.platform.commons.support.AnnotationSupport
-import org.junit.runners.model.*
 import java.lang.reflect.Method
-import java.util.concurrent.CountDownLatch
-import java.util.concurrent.ExecutionException
-import java.util.concurrent.FutureTask
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.TimeoutException
+
+public class CoroutinesTimeoutException(public val timeoutMs: Long): Exception("test timed out ofter $timeoutMs ms")
 
 public class CoroutinesTimeoutExtension: InvocationInterceptor {
 
@@ -107,7 +103,9 @@ public class CoroutinesTimeoutExtension: InvocationInterceptor {
         DebugProbes.enableCreationStackTraces = annotation.enableCoroutineCreationStackTraces
         DebugProbes.install()
         return try {
-            runWithTimeoutDumpingCoroutines(methodName, annotation.testTimeoutMs, annotation.cancelOnTimeout) {
+            runWithTimeoutDumpingCoroutines(methodName, annotation.testTimeoutMs, annotation.cancelOnTimeout,
+                { CoroutinesTimeoutException(annotation.testTimeoutMs) }
+            ) {
                 invocation.proceed()
             }
         } finally {

kotlinx-coroutines-debug/test/junit5/CoroutinesTimeoutTest.kt

@@ -4,6 +4,7 @@
 
 package junit5
 
+import kotlinx.coroutines.debug.junit5.*
 import org.assertj.core.api.*
 import org.junit.*
 import org.junit.platform.engine.discovery.DiscoverySelectors.*
@@ -36,10 +37,9 @@ private fun ListAssert<Event>.testFinishedSuccessfully(testName: String): ListAs
         finishedSuccessfully()
     ))
 
-private fun ListAssert<Event>.testTimedOut(testName: String, after: Int): ListAssert<Event> =
+private fun ListAssert<Event>.testTimedOut(testName: String, after: Long): ListAssert<Event> =
     haveExactly(1, event(
         test(testName),
-        finishedWithFailure(instanceOf(TestTimedOutException::class.java), message {
-            it.contains(Regex("\\b$after\\b"))
-        })
+        finishedWithFailure(Condition({ it is CoroutinesTimeoutException && it.timeoutMs == after },
+            "is CoroutinesTimeoutException($after)"))
     ))