GRPC Cleanup Extension
for JUnit 5


Topher Lamey|July 11, 2021

Learn how to use JUnit 5 and make test resources get automatically cleaned up to avoid test pollution.

The grpc-java project provides GrpcCleanupRule which is a helpful way to clean up GrpcService resources in JUnit 4 tests, but it doesn't work with JUnit 5. JUnit 5 doesn't support @Rule classes, but rather has moved that functionality into Extensions. Additionally, the GRPC folks have said they aren't moving to JUnit 5 for a while, and they suggest sticking with JUnit 4.

Here at StackHawk, we want to use JUnit 5 exclusively, and we want the test resources to get automatically cleaned up to avoid test pollution. So we wrote a JUnit 5 Extension called GrpcCleanupExtension that handles the channel/server resource cleanup after each test invocation.

The Problem

If you run a JUnit 5 test with theGrpcCleanupRule, you'll see something this in the logs:

2021-07-08 09:33:21,962 ERROR [Test worker] io.grpc.internal.ManagedChannelOrphanWrapper$ManagedChannelReference: *~*~*~ Channel ManagedChannelImpl{logId=348, target=directaddress:///6f89ffb5-fd55-46c3-94fa-1700e758525b} was not shutdown properly!!! ~*~*~*
    Make sure to call shutdown()/shutdownNow() and wait until awaitTermination() returns true.

Which means that the GrpcCleanupRule did not fire (because JUnit 5 doesn't support @Rule), and the previous server instance was replaced without shutting down correctly.

Grpc Cleanup Extension Use

Here's a look at how our Grpc Services are set up and how we use the cleanup Extension in tests.

Let's say we have a GrpcService implementation called IntegrationService that's set up like this:

class IntegrationService() : IntegrationServiceGrpc.IntegrationServiceImplBase() {
    // do integration stuff

When we go to write a test for IntegrationServer, we set up a GrpcCleanupExtension instance. The syntax is a little non-Kotlinesque because extensions need to be final static member variables (not in a companion object).

class IntegrationServiceTests {
    final val grpcCleanupExtension = GrpcCleanupExtension() // yes, the final is needed in kotlin

    lateinit var blockingStub: IntegrationServiceGrpc.IntegrationServiceBlockingStub

    fun setup() {
        val integrationService = IntegrationService()
        blockingStub = IntegrationServiceGrpc.newBlockingStub(grpcCleanupExtension.addService(integrationService))
    fun testThingOne() {

And that's it!

How The Cleanup Extension Works

Not surprisingly, the GrpcCleanupExtension looks very similar to the GrpcCleanupRule. It keeps track of a list of servers/channels, and shuts them down after each test execution in the afterEach function.

class GrpcCleanupExtension : AfterEachCallback {

    private var cleanupTargets: MutableList<CleanupTarget> = mutableListOf()

    companion object {
        private val logger = LoggerFactory.getLogger(

        const val TERMINATION_TIMEOUT_MS = 250L
        const val MAX_NUM_TERMINATIONS = 100

    fun addService(service: BindableService): ManagedChannel {
        val serverName: String = InProcessServerBuilder.generateName()


        val channel = InProcessChannelBuilder.forName(serverName)


        return channel

    override fun afterEach(context: ExtensionContext?) {
        cleanupTargets.forEach { cleanupTarget ->
            try {
                var count = 0
                do {
                    cleanupTarget.awaitTermination(TERMINATION_TIMEOUT_MS, TimeUnit.MILLISECONDS)
                    if (count > MAX_NUM_TERMINATIONS) {
                        logger.error("Hit max count $count trying to shut down down cleanupTarget $cleanupTarget")
                } while (!cleanupTarget.isTerminated())
            } catch (e: Exception) {
                logger.error("Problem shutting down cleanupTarget $cleanupTarget", e)

        if (isAllTerminated()) {
        } else {
            logger.error("Not all cleanupTargets are terminated")

    fun isAllTerminated(): Boolean = cleanupTargets.all { it.isTerminated() }

interface CleanupTarget {
    fun shutdown()
    fun awaitTermination(timeout: Long, timeUnit: TimeUnit): Boolean
    fun isTerminated(): Boolean

class ServerCleanupTarget(private val server: Server) : CleanupTarget {
    override fun shutdown() {

    override fun awaitTermination(timeout: Long, timeUnit: TimeUnit): Boolean =
        server.awaitTermination(timeout, timeUnit)

    override fun isTerminated(): Boolean = server.isTerminated

class ManagedChannelCleanupTarget(private val managedChannel: ManagedChannel) : CleanupTarget {
    override fun shutdown() {

    override fun awaitTermination(timeout: Long, timeUnit: TimeUnit): Boolean =
        managedChannel.awaitTermination(timeout, timeUnit)

    override fun isTerminated(): Boolean = managedChannel.isTerminated

C'est La Fin

In order to avoid JUnit 4 when testing GrpcService classes, we wrote a GrpcCleanupExtension. It hooks into the JUnit 5 test lifecycle to clean up server/channel resources. Enjoy!

Topher Lamey  |  July 11, 2021