From 4a1d56008436e90a7c38d9745613b1244bded05b Mon Sep 17 00:00:00 2001
From: Abhijit Sarkar <sarkar.abhijit@outlook.com>
Date: Fri, 14 Aug 2020 07:56:21 -0700
Subject: [PATCH] Closes #5: Extensions have to be stateless

---
 .../asarkar/grpc/test/GrpcCleanupExtension.kt | 68 +++++++++++++------
 1 file changed, 49 insertions(+), 19 deletions(-)

diff --git a/src/main/kotlin/com/asarkar/grpc/test/GrpcCleanupExtension.kt b/src/main/kotlin/com/asarkar/grpc/test/GrpcCleanupExtension.kt
index 179bf8e..2e17ace 100644
--- a/src/main/kotlin/com/asarkar/grpc/test/GrpcCleanupExtension.kt
+++ b/src/main/kotlin/com/asarkar/grpc/test/GrpcCleanupExtension.kt
@@ -30,8 +30,12 @@ import kotlin.reflect.KFunction1
  */
 class GrpcCleanupExtension :
     BeforeEachCallback, AfterEachCallback, ParameterResolver, BeforeAllCallback, AfterAllCallback {
-    private val resources = mutableMapOf<Boolean, MutableList<Resources>>()
-    private var resourcesField: Field? = null
+    companion object {
+        private val NAMESPACE: ExtensionContext.Namespace = ExtensionContext.Namespace
+            .create(*GrpcCleanupExtension::class.java.name.split(".").toTypedArray())
+        private const val RESOURCES = "resources"
+        private const val RESOURCES_FIELD = "resources-field"
+    }
 
     override fun beforeEach(ctx: ExtensionContext) {
         val resources = ctx.requiredTestMethod.parameters
@@ -41,28 +45,52 @@ class GrpcCleanupExtension :
         }
 
         if (shouldAccessResourcesField(ctx)) {
-            if (tryGet(ctx.requiredTestInstance) != null) {
+            if (tryGetField(ctx) != null) {
                 throw PreconditionViolationException(
                     "Either set lifecycle PER_CLASS or don't initialize Resources field"
                 )
             }
-            trySet(ctx.requiredTestInstance)
+            trySetField(ctx)
         }
     }
 
+    @Suppress("UNCHECKED_CAST")
+    private var ExtensionContext.resources: MutableMap<Boolean, MutableList<Resources>>
+        get() = getStore(NAMESPACE)
+            .getOrDefault(
+                RESOURCES,
+                MutableMap::class.java,
+                mutableMapOf<Boolean, MutableList<Resources>>()
+            ) as MutableMap<Boolean, MutableList<Resources>>
+        set(value) {
+            getStore(NAMESPACE)
+                .put(RESOURCES, value)
+        }
+
+    private var ExtensionContext.resourcesField: Field?
+        get() = getStore(NAMESPACE)
+            .get(
+                RESOURCES_FIELD,
+                Field::class.java
+            )
+        set(value) {
+            getStore(NAMESPACE)
+                .put(RESOURCES_FIELD, value)
+        }
+
     override fun afterEach(ctx: ExtensionContext) {
         var successful = true
-        this.resources[false]?.forEach {
+        ctx.resources[false]?.forEach {
             getCleanUpMethod(ctx)(it)
             successful = it.awaitReleased()
         }
 
         if (shouldAccessResourcesField(ctx)) {
-            tryGet(ctx.requiredTestInstance)?.also {
+            tryGetField(ctx)?.also {
                 getCleanUpMethod(ctx)(it)
                 successful = it.awaitReleased()
 
-                trySet(ctx.requiredTestInstance, null)
+                trySetField(ctx, null)
             }
         }
         if (!successful) throw PostconditionViolationException("One or more Resources couldn't be released")
@@ -79,7 +107,7 @@ class GrpcCleanupExtension :
                     parameterCtx.declaringExecutable.isAnnotationPresent(BeforeAll::class.java)
                 )
 
-        return Resources().also { resources.getOrPut(once, { mutableListOf() }).add(it) }
+        return Resources().also { extensionCtx.resources.getOrPut(once, { mutableListOf() }).add(it) }
     }
 
     override fun beforeAll(ctx: ExtensionContext) {
@@ -91,20 +119,20 @@ class GrpcCleanupExtension :
             } catch (e: ReflectiveOperationException) {
                 throw JUnitException("Illegal state: Cannot access Resources field", e)
             }
-            resourcesField = field
-            if (tryGet(ctx.testInstance.orElse(null)) == null) {
-                trySet(ctx.testInstance.orElse(null))
+            ctx.resourcesField = field
+            if (tryGetField(ctx) == null) {
+                trySetField(ctx)
             }
         }
     }
 
     override fun afterAll(ctx: ExtensionContext) {
         var successful = true
-        tryGet(ctx.testInstance.orElse(null))?.also {
+        tryGetField(ctx)?.also {
             getCleanUpMethod(ctx)(it)
             successful = it.awaitReleased()
         }
-        this.resources[true]?.forEach {
+        ctx.resources[true]?.forEach {
             getCleanUpMethod(ctx)(it)
             successful = it.awaitReleased()
         }
@@ -112,22 +140,24 @@ class GrpcCleanupExtension :
     }
 
     private fun shouldAccessResourcesField(ctx: ExtensionContext): Boolean {
-        return resourcesField != null &&
+        return ctx.resourcesField != null &&
             ctx.testInstanceLifecycle.orElse(null) != TestInstance.Lifecycle.PER_CLASS &&
-            !resourcesField.isStatic()
+            !ctx.resourcesField.isStatic()
     }
 
-    private fun trySet(target: Any?, value: Resources? = Resources()) {
+    private fun trySetField(ctx: ExtensionContext, value: Resources? = Resources()) {
         try {
-            resourcesField?.takeIf { target != null || it.isStatic() }?.set(target, value)
+            val target = ctx.testInstance.orElse(null)
+            ctx.resourcesField?.takeIf { target != null || it.isStatic() }?.set(target, value)
         } catch (e: ReflectiveOperationException) {
             throw JUnitException("Illegal state: Cannot set Resources field", e)
         }
     }
 
-    private fun tryGet(target: Any?): Resources? {
+    private fun tryGetField(ctx: ExtensionContext): Resources? {
         return try {
-            resourcesField?.takeIf { target != null || it.isStatic() }?.get(target) as Resources?
+            val target = ctx.testInstance.orElse(null)
+            ctx.resourcesField?.takeIf { target != null || it.isStatic() }?.get(target) as Resources?
         } catch (e: ReflectiveOperationException) {
             throw JUnitException("Illegal state: Cannot get Resources field", e)
         }