From 2d81750a13c84ef060ab46899793ed4e2b3ad522 Mon Sep 17 00:00:00 2001
From: Adrien Grand <jpountz@gmail.com>
Date: Mon, 26 Dec 2016 14:55:22 +0100
Subject: [PATCH] Make ESTestCase resilient to initialization errors.

---
 .../elasticsearch/common/logging/DeprecationLogger.java   | 3 ++-
 .../src/main/java/org/elasticsearch/test/ESTestCase.java  | 8 +++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/org/elasticsearch/common/logging/DeprecationLogger.java b/core/src/main/java/org/elasticsearch/common/logging/DeprecationLogger.java
index 55f89ce84ad..39b882f1a1d 100644
--- a/core/src/main/java/org/elasticsearch/common/logging/DeprecationLogger.java
+++ b/core/src/main/java/org/elasticsearch/common/logging/DeprecationLogger.java
@@ -25,6 +25,7 @@ import org.elasticsearch.common.SuppressLoggerChecks;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 
 import java.util.Iterator;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CopyOnWriteArraySet;
 
@@ -63,7 +64,7 @@ public class DeprecationLogger {
      * @throws IllegalStateException if this {@code threadContext} has already been set
      */
     public static void setThreadContext(ThreadContext threadContext) {
-        assert threadContext != null;
+        Objects.requireNonNull(threadContext, "Cannot register a null ThreadContext");
 
         // add returning false means it _did_ have it already
         if (THREAD_CONTEXT.add(threadContext) == false) {
diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
index 567c987ebf0..ca4f846fe39 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
@@ -255,6 +255,7 @@ public abstract class ESTestCase extends LuceneTestCase {
     @Before
     public final void before()  {
         logger.info("[{}]: before test", getTestName());
+        assertNull("Thread context initialized twice", threadContext);
         if (enableWarningsCheck()) {
             this.threadContext = new ThreadContext(Settings.EMPTY);
             DeprecationLogger.setThreadContext(threadContext);
@@ -272,8 +273,13 @@ public abstract class ESTestCase extends LuceneTestCase {
     @After
     public final void after() throws Exception {
         checkStaticState();
-        if (enableWarningsCheck()) {
+        // We check threadContext != null rather than enableWarningsCheck()
+        // because after methods are still called in the event that before
+        // methods failed, in which case threadContext might not have been
+        // initialized
+        if (threadContext != null) {
             ensureNoWarnings();
+            threadContext = null;
         }
         ensureAllSearchContextsReleased();
         ensureCheckIndexPassed();