From a103fee2b380f3ad075e558a2a21e6c209d2dfd7 Mon Sep 17 00:00:00 2001
From: arne <arne@apache.org>
Date: Fri, 26 Apr 2024 08:53:36 +0200
Subject: [PATCH] #7: Implement fallback for empty ExtensionContext

---
 src/main/java/rocks/limburg/cdimock/CdiMockExtension.java | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/main/java/rocks/limburg/cdimock/CdiMockExtension.java b/src/main/java/rocks/limburg/cdimock/CdiMockExtension.java
index 2411213..65c5e0d 100644
--- a/src/main/java/rocks/limburg/cdimock/CdiMockExtension.java
+++ b/src/main/java/rocks/limburg/cdimock/CdiMockExtension.java
@@ -53,6 +53,7 @@
 public class CdiMockExtension implements Extension {
 
     private static ThreadLocal<ExtensionContext> beforeAllContext = new ThreadLocal<ExtensionContext>();
+    private static ExtensionContext latestExtensionContext;
 
     private Set<Class<?>> excludedClasses;
     private Set<Class<? extends Annotation>> excludingAnnotationTypes;
@@ -60,14 +61,19 @@ public class CdiMockExtension implements Extension {
 
     public static void beforeAll(ExtensionContext context) {
         beforeAllContext.set(context);
+        latestExtensionContext = context;
     }
 
     public static void afterAll(ExtensionContext context) {
+        ExtensionContext beforeAllContext = beforeAllContext.get();
         beforeAllContext.remove();
+        if (beforeAllContext == latestExtensionContext) {
+            latestExtensionContext = null;
+        }
     }
 
     public Optional<ExtensionContext> getExtensionContext() {
-        return ofNullable(beforeAllContext.get());
+        return ofNullable(beforeAllContext.get()).or(() -> ofNullable(latestExtensionContext));
     }
 
     public void excludeClasses(@Observes ProcessAnnotatedType<?> event) {