commonMain.io.kotest.framework.concurrency.eventually.kt Maven / Gradle / Ivy

Go to download
package io.kotest.framework.concurrency

import io.kotest.assertions.ErrorCollectionMode
import io.kotest.assertions.errorCollector
import io.kotest.assertions.failure
import io.kotest.common.ExperimentalKotest
import io.kotest.mpp.timeInMillis
import kotlinx.coroutines.delay
import kotlin.reflect.KClass
import kotlin.time.Duration
import kotlin.time.ExperimentalTime

@OptIn(ExperimentalKotest::class)
typealias EventuallyStateFunction = (EventuallyState) -> U
typealias ThrowablePredicate = (Throwable) -> Boolean

@ExperimentalKotest
data class EventuallyConfig(
   val duration: Long = defaultDuration,
   val interval: Interval = defaultInterval,
   val initialDelay: Long = defaultDelay,
   val retries: Int = Int.MAX_VALUE,
   val suppressExceptions: Set> = setOf(),
   val suppressExceptionIf: ThrowablePredicate? = null,
   val listener: EventuallyStateFunction? = null,
   val predicate: EventuallyStateFunction? = null,
   val shortCircuit: EventuallyStateFunction? = null,
)

@ExperimentalKotest
private fun  EventuallyConfig.toBuilder() = EventuallyBuilder().apply {
   duration = [email protected]
   interval = [email protected]
   initialDelay = [email protected]
   retries = [email protected]
   suppressExceptions = [email protected]
   suppressExceptionIf = [email protected]
   listener = [email protected]
   predicate = [email protected]
   shortCircuit = [email protected]
}

@ExperimentalKotest
class EventuallyBuilder {
   var duration: Long = defaultDuration
   var interval: Interval = defaultInterval
   var initialDelay: Long = defaultDelay
   var retries: Int = Int.MAX_VALUE
   var suppressExceptions: Set> = setOf(AssertionError::class)
   var suppressExceptionIf: ThrowablePredicate? = null
   var listener: EventuallyStateFunction? = null
   var predicate: EventuallyStateFunction? = null
   var shortCircuit: EventuallyStateFunction? = null

   fun build() = EventuallyConfig(
      duration = duration, interval = interval, initialDelay = initialDelay, retries = retries,
      suppressExceptions = suppressExceptions, suppressExceptionIf = suppressExceptionIf,
      listener = listener, predicate = predicate, shortCircuit = shortCircuit
   )
}

@ExperimentalKotest
class EventuallyShortCircuitException(override val message: String) : Throwable()

@ExperimentalKotest
data class EventuallyState(
   val result: T?,
   val start: Long,
   val end: Long,
   val times: Int,
   val firstError: Throwable?,
   val thisError: Throwable?,
)

@ExperimentalKotest
private class EventuallyControl(val config: EventuallyConfig<*>) {
   val start = timeInMillis()
   val end = start + config.duration

   var times = 0
   var predicateFailedTimes = 0

   var firstError: Throwable? = null
   var lastError: Throwable? = null

   var lastDelayPeriod: Long = 0L
   var lastInterval: Long = 0L

   fun exceptionIsNotSuppressible(e: Throwable): Boolean {
      if (firstError == null) {
         firstError = e
      } else {
         lastError = e
      }

      if (EventuallyShortCircuitException::class.isInstance(e)) {
         return true
      }

      if (config.suppressExceptionIf?.invoke(e) == false) {
         return true
      }

      return !config.suppressExceptions.any { it.isInstance(e) }
   }

   fun  toState(result: T?) = EventuallyState(result = result, start = start, end = end, times = times, firstError = firstError, thisError = lastError)

   suspend fun step() {
      lastInterval = config.interval.next(++times)
      val delayMark = timeInMillis()
      delay(lastInterval)
      lastDelayPeriod = timeInMillis() - delayMark
   }

   fun attemptsRemaining() = timeInMillis() < end && times < config.retries

   /**
    * if we only executed once, and the last delay was > last interval, we didn't get a chance to run again so we run once more before exiting
    */
   fun isLongWait() = times == 1 && lastDelayPeriod > lastInterval

   fun buildFailureMessage() = StringBuilder().apply {
      appendLine("Eventually block failed after ${config.duration}ms; attempted $times time(s); ${config.interval} delay between attempts")

      if (predicateFailedTimes > 0) {
         appendLine("The provided predicate failed $predicateFailedTimes times")
      }

      firstError?.run {
         appendLine("The first error was caused by: ${this.message}")
         appendLine(this.stackTraceToString())
      }

      lastError?.run {
         appendLine("The last error was caused by: ${this.message}")
         appendLine(this.stackTraceToString())
      }
   }.toString()
}

@ExperimentalKotest
suspend operator fun  EventuallyConfig.invoke(f: suspend () -> T): T {
   delay(initialDelay)

   val originalAssertionMode = errorCollector.getCollectionMode()
   errorCollector.setCollectionMode(ErrorCollectionMode.Hard)

   val control = EventuallyControl(this)

   try {
      while (control.attemptsRemaining() || control.isLongWait()) {
         try {
            val result = f()
            val state = control.toState(result)

            listener?.invoke(state)

            when (shortCircuit?.invoke(state)) {
               null, false -> Unit
               true -> throw EventuallyShortCircuitException("The provided shortCircuit function caused eventually to exit early: $state")
            }

            when (predicate?.invoke(state)) {
               null, true -> return result
               false -> control.predicateFailedTimes++
            }
         } catch (e: Throwable) {
            val notSuppressible = control.exceptionIsNotSuppressible(e)
            listener?.invoke(control.toState(null))
            if (notSuppressible) {
               throw e
            }
         }

         control.step()
      }
   } finally {
      errorCollector.setCollectionMode(originalAssertionMode)
   }

   throw failure(control.buildFailureMessage())
}

// region eventually

@ExperimentalKotest
suspend fun  eventually(
   config: EventuallyConfig,
   configure: EventuallyBuilder.() -> Unit,
   @BuilderInference test: suspend () -> T
): T {
   val resolvedConfig = config.toBuilder().apply(configure).build()
   return resolvedConfig.invoke(test)
}

@ExperimentalKotest
suspend fun  eventually(
   configure: EventuallyBuilder.() -> Unit, @BuilderInference test: suspend () -> T
): T {
   val config = EventuallyBuilder().apply(configure).build()
   return config.invoke(test)
}

@ExperimentalKotest
suspend fun  eventually(
   config: EventuallyConfig, @BuilderInference test: suspend () -> T
): T {
   return config.invoke(test)
}

@ExperimentalTime
@ExperimentalKotest
suspend fun  eventually(duration: Duration, test: suspend () -> T): T =
   eventually(duration.inWholeMilliseconds, test)

@ExperimentalKotest
suspend fun  eventually(duration: Long, test: suspend () -> T): T = eventually({ this.duration = duration }, test)

// endregion

// region until

@ExperimentalKotest
suspend fun until(
   config: EventuallyConfig, configure: EventuallyBuilder.() -> Unit, @BuilderInference test: suspend () -> Boolean
) {
   val builder = config.toBuilder()
   builder.predicate = { it.result == true }
   builder.apply(configure)
   builder.build().invoke(test)
}

@ExperimentalKotest
suspend fun until(
   configure: EventuallyBuilder.() -> Unit, @BuilderInference test: suspend () -> Boolean
) {
   val builder = EventuallyBuilder()
   builder.predicate = { it.result == true }
   builder.apply(configure)
   builder.build().invoke(test)
}

@ExperimentalTime
@ExperimentalKotest
suspend fun until(duration: Duration, test: suspend () -> Boolean) = until(millis = duration.inWholeMilliseconds, test)

@ExperimentalKotest
suspend fun until(millis: Long, test: suspend () -> Boolean) = until({ this.duration = millis }, test)

// endregion