Merge pull request #621 from jglick/InboundAgentRule-withLogger

`InboundAgentRule.Options.Builder.withLogger`
jenkinsci · Jul 17, 2023 · 787a641 · 787a641
2 parents a7371de + 39c760a
commit 787a641
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 14 deletions.
diff --git a/src/main/java/org/jvnet/hudson/test/InboundAgentRule.java b/src/main/java/org/jvnet/hudson/test/InboundAgentRule.java
@@ -31,6 +31,7 @@
 import hudson.model.Descriptor;
 import hudson.model.Node;
 import hudson.model.Slave;
+import hudson.remoting.VirtualChannel;
 import hudson.slaves.DumbSlave;
 import hudson.slaves.JNLPLauncher;
 import hudson.slaves.RetentionStrategy;
@@ -40,10 +41,13 @@
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.apache.commons.io.FileUtils;
 import org.apache.tools.ant.util.JavaEnvUtils;
@@ -84,7 +88,7 @@ public static final class Options implements Serializable {
         private boolean webSocket;
         @CheckForNull private String tunnel;
         private boolean start = true;
-
+        private final Map<String, Level> loggers = new LinkedHashMap<>();
         private String label;
         private final PrefixedOutputStream.Builder prefixedOutputStreamBuilder = PrefixedOutputStream.builder();
 
@@ -206,6 +210,20 @@ public Builder label(String label) {
                 return this;
             }
 
+            public Builder withLogger(Class<?> clazz, Level level) {
+                return withLogger(clazz.getName(), level);
+            }
+
+            public Builder withPackageLogger(Class<?> clazz, Level level) {
+                return withLogger(clazz.getPackageName(), level);
+            }
+
+            public Builder withLogger(String logger, Level level) {
+                options.loggers.put(logger, level);
+                return this;
+            }
+
+
             /**
              * Build and return an {@link Options}.
              *
@@ -270,7 +288,7 @@ public void start(@NonNull JenkinsRule r, Options options) throws Exception {
         Objects.requireNonNull(name);
         stop(r, name);
         start(GetAgentArguments.get(r, name), options);
-        WaitForAgentOnline.wait(r, name);
+        WaitForAgentOnline.wait(r, name, options.loggers);
     }
 
     /**
@@ -281,7 +299,7 @@ public void start(@NonNull RealJenkinsRule r, Options options) throws Throwable
         Objects.requireNonNull(name);
         stop(r, name);
         start(r.runRemotely(new GetAgentArguments(name)), options);
-        r.runRemotely(new WaitForAgentOnline(name));
+        r.runRemotely(new WaitForAgentOnline(name, options.loggers));
     }
 
     @SuppressFBWarnings(value = {"COMMAND_INJECTION", "NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE"}, justification = "just for test code")
@@ -432,17 +450,19 @@ private static AgentArguments get(JenkinsRule r, String name) throws IOException
 
     private static class WaitForAgentOnline implements RealJenkinsRule.Step {
         private final String name;
+        private final Map<String, Level> loggers;
 
-        WaitForAgentOnline(String name) {
+        WaitForAgentOnline(String name, Map<String, Level> loggers) {
             this.name = name;
+            this.loggers = loggers;
         }
 
         @Override
         public void run(JenkinsRule r) throws Throwable {
-            wait(r, name);
+            wait(r, name, loggers);
         }
 
-        static void wait(JenkinsRule r, String name) throws Exception {
+        static void wait(JenkinsRule r, String name, Map<String, Level> loggers) throws Exception {
             Node node = r.jenkins.getNode(name);
             if (node == null) {
                 throw new AssertionError("no such agent: " + name);
@@ -451,6 +471,11 @@ static void wait(JenkinsRule r, String name) throws Exception {
                 throw new AssertionError("agent is not a Slave: " + name);
             }
             r.waitOnline((Slave) node);
+            if (!loggers.isEmpty()) {
+                VirtualChannel channel = node.getChannel();
+                assert channel != null;
+                channel.call(new JenkinsRule.RemoteLogDumper(null, loggers, false));
+            }
         }
     }
 

diff --git a/src/main/java/org/jvnet/hudson/test/JenkinsRule.java b/src/main/java/org/jvnet/hudson/test/JenkinsRule.java
@@ -108,6 +108,7 @@
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.OutputStream;
+import java.io.PrintStream;
 import java.io.UncheckedIOException;
 import java.lang.annotation.Annotation;
 import java.lang.management.ThreadInfo;
@@ -1193,27 +1194,29 @@ public void showAgentLogs(Slave s, LoggerRule loggerRule) throws Exception {
      * @param loggers {@link Logger#getName} tied to log level
      */
     public void showAgentLogs(Slave s, Map<String, Level> loggers) throws Exception {
-        s.getChannel().call(new RemoteLogDumper(s.getNodeName(), loggers));
+        s.getChannel().call(new RemoteLogDumper(s.getNodeName(), loggers, true));
     }
 
-    private static final class RemoteLogDumper extends MasterToSlaveCallable<Void, RuntimeException> {
+    static final class RemoteLogDumper extends MasterToSlaveCallable<Void, RuntimeException> {
         private final String name;
         private final Map<String, Level> loggers;
-        private final TaskListener stderr = StreamTaskListener.fromStderr();
+        private final TaskListener stderr;
         private final long start = DeltaSupportLogFormatter.start;
         @SuppressWarnings("MismatchedQueryAndUpdateOfCollection")
         private static final List<Logger> loggerReferences = new LinkedList<>();
-        RemoteLogDumper(String name, Map<String, Level> loggers) {
+        RemoteLogDumper(String name, Map<String, Level> loggers, boolean forward) {
             this.name = name;
             this.loggers = loggers;
+            stderr = forward ? StreamTaskListener.fromStderr() : null;
         }
         @Override public Void call() throws RuntimeException {
+            PrintStream ps = stderr != null ? stderr.getLogger() : System.err;
             Handler handler = new Handler() {
                 final Formatter formatter = new DeltaSupportLogFormatter();
                 @Override public void publish(LogRecord record) {
                     if (isLoggable(record)) {
-                        stderr.getLogger().print(formatter.format(record).replaceAll("(?m)^([ 0-9.]*)", "$1[" + name + "] "));
-                        stderr.getLogger().flush();
+                        ps.print(formatter.format(record).replaceAll("(?m)^([ 0-9.]*)", name != null ? "$1[" + name + "] " : "$1 "));
+                        ps.flush();
                     }
                 }
                 @Override public void flush() {}
@@ -1227,8 +1230,12 @@ private static final class RemoteLogDumper extends MasterToSlaveCallable<Void, R
                 loggerReferences.add(logger);
             });
             DeltaSupportLogFormatter.start = start; // match clock time on master
-            stderr.getLogger().println("Set up log dumper on " + name + ": " + loggers);
-            stderr.getLogger().flush();
+            if (name != null) {
+                ps.println("Set up log dumper on " + name + ": " + loggers);
+            } else {
+                ps.println("Set up log dumper: " + loggers);
+            }
+            ps.flush();
             return null;
         }
     }