Coroutines

Eliminating Coroutine leaks in tests

Rick Busarow

• 8 min read
Eliminating Coroutine leaks in tests

Coroutines are all the rage right now, and if you’re writing coroutines, you’re probably writing tests for them. Writing tests for coroutines can be difficult, not only because of concurrency concerns, but because it is very easy to create leaks which live on past the lifecycle of an individual test.

Consider this example class:

class Subject {

  var someBoolean = false

  fun CoroutineScope.loop() {

    someBoolean = true
    
    launch {
      repeat(10) { count ->
        delay(timeMillis = 1_000)
        println("loop is running -- $count")
      }
      println("all done")
    }
  }
}

This single function loop() will immediately update someBoolean, then create a new coroutine with launch { }. This coroutine will run asynchronously for 10 seconds, printing a reminder that it’s running every second. Since launch { } is non-blocking, loop() will return before the new coroutine finishes.

Now, let’s consider a simple JUnit 4 test class which may be testing that loop function.

import io.kotlintest.shouldBe
class MyLeakyTest {

  val scope = CoroutineScope(Dispatchers.Unconfined)

  val subject = Subject()

  @Before
  fun before() {
    println("before my leaky test")
  }

  @After
  fun after() {
    println("after my leaky test")
  }

  @Test
  fun `create a leak`() {

    with(subject) {
      scope.loop()
    }

    subject.someBoolean shouldBe true
    println("my leaky test has completed")
  }

}

If we run this test class, the order of execution will be:

  1. before()
  2. create a leak()
  3. after()

Here’s the proof:

Everything seems to be fine, right? Well, that’s because we’re only executing this one test, and the entire JVM is shut down immediately after. But of course, when we run tests, we want to run all of our tests. Let’s simulate that by adding one more test class:

class SomeOtherTest {

  @Test
  fun `another test`() {

    println("some other tests would run now")

    runBlocking { delay(11_000) }
  }
}

Because this test uses runBlocking { }, the delay it calls will ensure that the test doesn’t complete for 11 seconds. This mimics the execution of other tests while a leak is happening.

If we run both test classes together, we can see that the side effects of the original test will live on long after the test itself has completed:

Now this single leak is pretty harmless. It’s just spamming the console, and it terminates on its own. But what if it was manipulating some shared state? That would pollute other tests, making them non-hermetic. Or, what if it was actually making a network call, and the mock it was intended to use was destroyed in @After? Would it cause an exception? Would it just suspend indefinitely waiting for a response?

Stopping the leak

Fortunately, there’s a simple fix to this leak. We simply need to cancel the CoroutineScope after execution of any test, like so:

@After
fun after() {
  scope.cancel()
  println("after my leaky test")
}

This ensures that any Jobs created by that scope will stop immediately, thereby stopping any leaks. If we re-run the test, we can verify that:

Note that the second test still ran for just over 11 seconds, but we never heard from the loop() function again. Yay!

The problem is… People make mistakes. That’s why we have tests. And our solution to preventing leaks now hinges upon every developer remembering to call scope.cancel() in the teardown of every test. I don’t trust myself to do that, and I don’t want to put the burden of remembering this crucial step on my fellow developers. Of course, we can come up with a better way.

Rules!

In JUnit 4, our better way comes in the form of org.junit.Rule. Rules are aware of the lifecycle of the tests, and are designed to automatically perform tasks before and/or after the test (Statement) is executed. Here is an example:

class CoroutineTestRule : TestRule, CoroutineScope {
  
  override val coroutineContext: CoroutineContext
    = Job() + Dispatchers.Unconfined

  override fun apply(
    base: Statement, description: Description?
  ) = object : Statement() {
    override fun evaluate() {
      println("before base.evaluate()")
      base.evaluate()
      println("after base.evaluate()")
      this@CoroutineTestRule.cancel() // cancels CoroutineScope
    }
  }
}

We make the Rule implement CoroutineScope, so it can double as the CoroutineScope in our tests. Execution of the actual test happens during base.eveluate(), then @After is called, and then the rule will cancel its own scope.

We can then utilize this rule in our test like this:

import io.kotlintest.shouldBe
class MyFormerlyLeakyTest {

  @get:Rule val scope = CoroutineTestRule()

  val subject = Subject()

  @Before
  fun before() {
    println("before my formerly leaky test")
  }

  @After
  fun after() {
    println("after my formerly leaky test")
  }

  @Test
  fun `create a potential leak`() {

    with(subject) {
      scope.loop()
    }

    subject.someBoolean shouldBe true
  }
}

Note that in order for Kotlin to understand the Rule, we need to add an @get:Rule annotation to the property.

Also note that we can continue to utilize this rule as a CoroutineScope.

Finally, note that we don’t need to put anything in @After. If it weren’t for my print statements, we could omit @Before and @After entirely. But since we have the print statements, we’d better use them. Here’s what we see if we run both our tests:

So we have no leaks still, and the cancellation of the CoroutineScope is handled for us, and for all other coroutine tests in the future. We can now write lots of safe tests with minimal boilerplate:

import io.kotlintest.shouldBe
class MyFormerlyLeakyTest {

  @get:Rule val scope = CoroutineTestRule()

  val subject = Subject()

  @Test
  fun `create a potential leak`() {

    with(subject) {
      scope.loop()
    }

    subject.someBoolean shouldBe true
  }
}
class AnotherTest {

  @get:Rule val scope = CoroutineTestRule()
  @Test
  fun `assume we're Unconfined`() {

    var updated = false
    scope.launch { updated = true }

    updated shouldBe true
  }
}

Building upon the rule

Now that we have a basic Rule which we can utilize, we can enhance its utility with the new testing tools introduced in kotlinx.coroutines-test v1.2.1. You can read all about the awesome new functionality here. You should also be sure to follow the author of that code,

Sean McQuillan, because he’s bound to cover testing in his seriesCoroutines on Android.

So that this post doesn’t get too long, I’m going to resist the urge to extol the virtues of this library any more than I need to. I’m only going to mention its features as they relate to the Rule.

So, the first step would be to modify our CoroutineTestRule so that it implements TestCoroutineScope instead of the normal CoroutineScope. For this, we’ll use Kotlin’s class delegation and the TestCoroutineScope() factory function:

class CoroutineTestRule : TestRule, 
                        TestCoroutineScope by TestCoroutineScope() {
  override fun apply(
    base: Statement, description: Description?
  ) = object : Statement() {
    override fun evaluate() {
      @Throws(Throwable::class)
      override fun evaluate() {

        base.evaluate()

        this@CoroutineTestRule.cancel()
      }
    }
  }
}

Next, one common issue with testing coroutines is use of Dispatchers.Main. We don’t have access to a real “main” thread in JVM tests, so we need to use the Dispatchers.setMain() testing utility function and pass it a CoroutineDispatcher. We can do this by extracting the dispatcher from the TestCoroutineScope like this:

class CoroutineTestRule : TestRule,
                       TestCoroutineScope by TestCoroutineScope() {
  val dispatcher =
    coroutineContext[ContinuationInterceptor]
        as TestCoroutineDispatcher
  override fun apply(
    base: Statement, description: Description?
  ) = object : Statement() {
    override fun evaluate() {
      @Throws(Throwable::class)
      override fun evaluate() {

        base.evaluate()

        this@CoroutineTestRule.cancel()
      }
    }
  }
}

And then setting and resetting Main in our rule like this:

class CoroutineTestRule : TestRule,
                       TestCoroutineScope by TestCoroutineScope() {
  val dispatcher =
    coroutineContext[ContinuationInterceptor]
        as TestCoroutineDispatcher
  override fun apply(
    base: Statement, description: Description?
  ) = object : Statement() {
    override fun evaluate() {
      @Throws(Throwable::class)
      override fun evaluate() {
        Dispatchers.setMain(dispatcher)

        base.evaluate()
        this@CoroutineTestRule.cancel()
        Dispatchers.resetMain()
      }
    }
  }
}

Finally, we can do away with the normal cancellation of the scope, because TestCoroutineScope has something better. The cleanupTestCoroutines()function will actually immediately throw an UncompletedCoroutinesError(and thereby fail your tests) if there are any coroutines still active. Suffice it to say this is a really good move, so we can add it like this:

class CoroutineTestRule : TestRule,
                      TestCoroutineScope by TestCoroutineScope() {

  val dispatcher =
    coroutineContext[ContinuationInterceptor]
        as TestCoroutineDispatcher
  
  override fun apply(
    base: Statement, description: Description?
  ) = object : Statement() {
    override fun evaluate() {
      @Throws(Throwable::class)
      override fun evaluate() {
        Dispatchers.setMain(dispatcher)

        base.evaluate()

        cleanupTestCoroutines()
        Dispatchers.resetMain()
      }
    }
  }
}

Trying our new Rule

With these changes, if we run our old leaky test everything should be great, right? Actually no — we get an exception:

This is intended. If this was testing code we actually care about, we would now have to go back and refactor our code or tests to properly handle that coroutine. For example, what if after several seconds, that unfinished coroutine would cause an Exception in a production workflow? Simply cancelling the scope would allow for the test to pass despite that defect. With this new functionality, we’re forced to account for all leaks instead of simply ignoring them. But this isn’t a post about how awesome this library is…

JUnit 5

Now that we have a JUnit 4 rule doing everything we want, what about JUnit 5? JUnit 5 doesn’t support Rules. Instead, it utilizes Extensions, and unlike a Rule, an Extension isn’t a property which can be accessed from the test class. Extensions instead have to inject properties to a class — through reflection, or through a little creative casting. We will utilize the second option since it’s cleaner and more performant.

Note — if you’ve been following along with actual code of your own, be sure to switch your @Test imports from org.junit.Test to org.junit.jupiter.api.Test or the Extension won’t get a chance to work.

First, let’s create the properties and handle their lifecycle. Since the Extension isn’t a property of the class, there’s no sense in making it implement a TestCoroutineScope interface. This implementation is therefore a little bit more straightforward:

@ExperimentalCoroutinesApi
class TestCoroutineExtension : BeforeAllCallback, 
                               AfterEachCallback, 
                               AfterAllCallback {
  val dispatcher = TestCoroutineDispatcher()
  val testScope = TestCoroutineScope(dispatcher)

  override fun beforeAll(context: ExtensionContext?) {
    Dispatchers.setMain(dispatcher)
  }

  override fun afterEach(context: ExtensionContext?) {
    testScope.cleanupTestCoroutines()
  }

  override fun afterAll(context: ExtensionContext?) {
    Dispatchers.resetMain()
  }
}

We can now use one of my favorite tricks for utilizing an Extension. Instead of annotating the test class, we can annotate an interface, then have the test class implement the interface:

@ExtendWith(TestCoroutineExtension::class)
interface CoroutineTest

class MyFormerlyLeakyTest : CoroutineTest {

  ...
}
class AnotherTest : CoroutineTest {

  ...
}

But this extension has properties to share, right? How do we do that? First, let’s add the properties to the interface:

@ExtendWith(TestCoroutineExtension::class)
interface CoroutineTest {
  var testScope: TestCoroutineScope
  var dispatcher: TestCoroutineDispatcher
}

Now we can add a new interface to the extension — TestInstancePostProcessor. This is executed after initialization and before beforeAll.

@ExperimentalCoroutinesApi
class TestCoroutineExtension : TestInstancePostProcessor,
                               BeforeAllCallback,
                               AfterEachCallback,
                               AfterAllCallback {
 
  val dispatcher = TestCoroutineDispatcher()
  val testScope = TestCoroutineScope(dispatcher)
 
  override fun postProcessTestInstance(
    testInstance: Any?, context: ExtensionContext?
  ) {
    
  }

  override fun beforeAll(context: ExtensionContext?) {
    Dispatchers.setMain(dispatcher)
  }

  override fun afterEach(context: ExtensionContext?) {
    testScope.cleanupTestCoroutines()
  }

  override fun afterAll(context: ExtensionContext?) {
    Dispatchers.resetMain()
  }
}

And now, we can cast testInstance to our CoroutineTest interface, then access its properties and set them!

@ExperimentalCoroutinesApi
class TestCoroutineExtension : TestInstancePostProcessor,
                               BeforeAllCallback,
                               AfterEachCallback,
                               AfterAllCallback {
 
  val dispatcher = TestCoroutineDispatcher()
  val testScope = TestCoroutineScope(dispatcher)
 
  override fun postProcessTestInstance(
    testInstance: Any?, context: ExtensionContext?
  ) {
    testInstance as CoroutineTest
    testInstance.testScope = testScope
    testInstance.dispatcher = dispatcher
  }

  override fun beforeAll(context: ExtensionContext?) {
    Dispatchers.setMain(dispatcher)
  }

  override fun afterEach(context: ExtensionContext?) {
    testScope.cleanupTestCoroutines()
  }

  override fun afterAll(context: ExtensionContext?) {
    Dispatchers.resetMain()
  }
}

Now, our test class must be updated to implement the new properties:

class MyFormerlyLeakyTest : CoroutineTest {

  override lateinit var testScope: TestCoroutineScope
  override lateinit var dispatcher: TestCoroutineDispatcher

  val subject = Subject()

  @Test
  fun `create a potential leak`() {

    with(subject) {
      testScope.loop()
    }

    assertTrue(subject.someBoolean)
  }
}

Now, the Extension (hidden in the interface) will automatically initialize the properties, then handle their complete lifecycle for us. Yay!

Conclusion

We demonstrated how a coroutine can leak in tests, and how to stop that bleeding with scope.cancel(). We then created a Rule to abstract that logic, and improved that rule with kotlinx.coroutines-test. Then, we implemented the same logic in a JUnit 5 Extension.

Thank you for reading! I hope that you found something useful!

Source Code

You can find the complete source for both the Rule and Extension in a Github gist here: https://bit.ly/2HyPn4j

Sign up for more like this.

Enter your email
Subscribe