From 581eb618b37c94cbd9f62eca497a628db4f39f7f Mon Sep 17 00:00:00 2001
From: Simon Willnauer <simonw@apache.org>
Date: Mon, 9 Dec 2013 11:02:13 +0100
Subject: [PATCH] Add Class and Package Level support for TestLogging

---
 .../test/junit/annotations/TestLogging.java   |  7 +-
 .../test/junit/listeners/LoggingListener.java | 69 +++++++++++++------
 2 files changed, 52 insertions(+), 24 deletions(-)

diff --git a/src/test/java/org/elasticsearch/test/junit/annotations/TestLogging.java b/src/test/java/org/elasticsearch/test/junit/annotations/TestLogging.java
index f56bc6858ba..82bc5cb9eeb 100644
--- a/src/test/java/org/elasticsearch/test/junit/annotations/TestLogging.java
+++ b/src/test/java/org/elasticsearch/test/junit/annotations/TestLogging.java
@@ -18,11 +18,14 @@
 
 package org.elasticsearch.test.junit.annotations;
 
-import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PACKAGE;
+import static java.lang.annotation.ElementType.TYPE;
+
 /**
  * Annotation used to set a custom log level for a specific test method.
  *
@@ -32,7 +35,7 @@ import java.lang.annotation.Target;
  * or just @TestLogging("_root:DEBUG,cluster.metadata:TRACE") since we start the test with -Des.logger.prefix=
  */
 @Retention(RetentionPolicy.RUNTIME)
-@Target(ElementType.METHOD)
+@Target({PACKAGE, TYPE, METHOD})
 public @interface TestLogging {
     String value();
 }
diff --git a/src/test/java/org/elasticsearch/test/junit/listeners/LoggingListener.java b/src/test/java/org/elasticsearch/test/junit/listeners/LoggingListener.java
index 10456280a0c..64d5fe3e54d 100644
--- a/src/test/java/org/elasticsearch/test/junit/listeners/LoggingListener.java
+++ b/src/test/java/org/elasticsearch/test/junit/listeners/LoggingListener.java
@@ -23,6 +23,7 @@ import org.elasticsearch.common.logging.ESLoggerFactory;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.junit.runner.Description;
+import org.junit.runner.Result;
 import org.junit.runner.notification.RunListener;
 
 import java.util.HashMap;
@@ -41,35 +42,30 @@ import java.util.Map;
 public class LoggingListener extends RunListener {
 
     private Map<String, String> previousLoggingMap;
+    private Map<String, String> previousClassLoggingMap;
+    private Map<String, String> previousPackageLoggingMap;
+
+    @Override
+    public void testRunStarted(Description description) throws Exception {
+        previousPackageLoggingMap = processTestLogging( description.getTestClass().getPackage().getAnnotation(TestLogging.class));
+        previousClassLoggingMap = processTestLogging(description.getAnnotation(TestLogging.class));
+    }
+
+    @Override
+    public void testRunFinished(Result result) throws Exception {
+        previousClassLoggingMap = reset(previousClassLoggingMap);
+        previousPackageLoggingMap = reset(previousPackageLoggingMap);
+    }
 
     @Override
     public void testStarted(Description description) throws Exception {
-        TestLogging testLogging = description.getAnnotation(TestLogging.class);
-        if (testLogging != null) {
-            this.previousLoggingMap = new HashMap<String, String>();
-            String[] loggersAndLevels = testLogging.value().split(",");
-            for (String loggerAndLevel : loggersAndLevels) {
-                String[] loggerAndLevelArray = loggerAndLevel.split(":");
-                if (loggerAndLevelArray.length >=2) {
-                    String loggerName = loggerAndLevelArray[0];
-                    String level = loggerAndLevelArray[1];
-                    ESLogger esLogger = resolveLogger(loggerName);
-                    this.previousLoggingMap.put(loggerName, esLogger.getLevel());
-                    esLogger.setLevel(level);
-                }
-            }
-        }
+        final TestLogging testLogging = description.getAnnotation(TestLogging.class);
+        previousLoggingMap = processTestLogging(testLogging);
     }
 
     @Override
     public void testFinished(Description description) throws Exception {
-        if (this.previousLoggingMap != null) {
-            for (Map.Entry<String, String> previousLogger : previousLoggingMap.entrySet()) {
-                ESLogger esLogger = resolveLogger(previousLogger.getKey());
-                esLogger.setLevel(previousLogger.getValue());
-            }
-            this.previousLoggingMap = null;
-        }
+        previousLoggingMap = reset(previousLoggingMap);
     }
 
     private static ESLogger resolveLogger(String loggerName) {
@@ -78,4 +74,33 @@ public class LoggingListener extends RunListener {
         }
         return Loggers.getLogger(loggerName);
     }
+
+    private Map<String, String> processTestLogging(TestLogging testLogging) {
+        if (testLogging == null) {
+            return null;
+        }
+        Map<String, String> map = new HashMap<String, String>();
+        final String[] loggersAndLevels = testLogging.value().split(",");
+        for (String loggerAndLevel : loggersAndLevels) {
+            String[] loggerAndLevelArray = loggerAndLevel.split(":");
+            if (loggerAndLevelArray.length >=2) {
+                String loggerName = loggerAndLevelArray[0];
+                String level = loggerAndLevelArray[1];
+                ESLogger esLogger = resolveLogger(loggerName);
+                map.put(loggerName, esLogger.getLevel());
+                esLogger.setLevel(level);
+            }
+        }
+        return map;
+    }
+
+    private Map<String, String> reset(Map<String, String> map) {
+        if (map != null) {
+            for (Map.Entry<String, String> previousLogger : map.entrySet()) {
+                ESLogger esLogger = resolveLogger(previousLogger.getKey());
+                esLogger.setLevel(previousLogger.getValue());
+            }
+        }
+        return null;
+    }
 }