Hamburger Icon

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

Read More

Add AppSec to Your CircleCI Pipeline With the StackHawk Orb

Add AppSec to Your CircleCI Pipeline With the StackHawk Orb

Application Security is Broken. Here is How We Intend to Fix It.

Application Security is Broken. Here is How We Intend to Fix It.

Using StackHawk in GitLab Know Before You Go (Live)

Using StackHawk in GitLab Know Before You Go (Live)