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:
@GRpcService
@Component
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).
@ExtendWith(SpringExtension::class)
class IntegrationServiceTests {
@JvmField
@RegisterExtension
final val grpcCleanupExtension = GrpcCleanupExtension() // yes, the final is needed in kotlin
lateinit var blockingStub: IntegrationServiceGrpc.IntegrationServiceBlockingStub
@BeforeEach
fun setup() {
val integrationService = IntegrationService()
blockingStub = IntegrationServiceGrpc.newBlockingStub(grpcCleanupExtension.addService(integrationService))
}
@Test
fun testThingOne() {
assertNotNull(blockingStub.thingOne())
}
}
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(GrpcCleanupExtension::class.java)
const val TERMINATION_TIMEOUT_MS = 250L
const val MAX_NUM_TERMINATIONS = 100
}
fun addService(service: BindableService): ManagedChannel {
val serverName: String = InProcessServerBuilder.generateName()
cleanupTargets.add(
ServerCleanupTarget(
InProcessServerBuilder
.forName(serverName)
.directExecutor()
.intercept(GlobalGrpcExceptionHandler())
.addService(service)
.build()
.start()
)
)
val channel = InProcessChannelBuilder.forName(serverName)
.directExecutor()
.build()
cleanupTargets.add(ManagedChannelCleanupTarget(channel))
return channel
}
override fun afterEach(context: ExtensionContext?) {
cleanupTargets.forEach { cleanupTarget ->
try {
var count = 0
cleanupTarget.shutdown()
do {
cleanupTarget.awaitTermination(TERMINATION_TIMEOUT_MS, TimeUnit.MILLISECONDS)
count++
if (count > MAX_NUM_TERMINATIONS) {
logger.error("Hit max count $count trying to shut down down cleanupTarget $cleanupTarget")
break
}
} while (!cleanupTarget.isTerminated())
} catch (e: Exception) {
logger.error("Problem shutting down cleanupTarget $cleanupTarget", e)
}
}
if (isAllTerminated()) {
cleanupTargets.clear()
} 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() {
server.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() {
managedChannel.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!
