From 7bc692eccb1729608108c101362bb908829544fd Mon Sep 17 00:00:00 2001
From: "Taro L. Saito" <leo@xerial.org>
Date: Mon, 6 May 2024 10:46:29 -0700
Subject: [PATCH] airspec (fix): Fixes #3523 and #3516 to properly set and
 reset log levels with -l,-L options (#3528)

Fixes #3523, #3516
---
 .../airspec/runner/AirSpecSbtRunner.scala     | 26 ++++++++++++++-----
 .../airspec/runner/AirSpecTaskRunner.scala    | 18 ++++++++++---
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/airspec/src/main/scala/wvlet/airspec/runner/AirSpecSbtRunner.scala b/airspec/src/main/scala/wvlet/airspec/runner/AirSpecSbtRunner.scala
index 88fe347baf..fa67d1bc63 100644
--- a/airspec/src/main/scala/wvlet/airspec/runner/AirSpecSbtRunner.scala
+++ b/airspec/src/main/scala/wvlet/airspec/runner/AirSpecSbtRunner.scala
@@ -59,9 +59,10 @@ private[airspec] object AirSpecSbtRunner extends LogSupport {
   def newRunner(args: Array[String], remoteArgs: Array[String], testClassLoader: ClassLoader): AirSpecSbtRunner = {
 
     // Set log level with -l (log level)
-    val remaining          = Array.newBuilder[String]
-    var i                  = 0
-    var logLevel: LogLevel = Logger.getDefaultLogLevel
+    val remaining           = Array.newBuilder[String]
+    var i                   = 0
+    var logLevel: LogLevel  = Logger.getDefaultLogLevel
+    val additionalLogLevels = Map.newBuilder[String, LogLevel]
     while (i < args.length) {
       args(i) match {
         case "-l" if i < args.length - 1 =>
@@ -71,8 +72,7 @@ private[airspec] object AirSpecSbtRunner extends LogSupport {
         case arg if arg.startsWith("-L") =>
           arg.stripPrefix("-L").split("=") match {
             case Array(pkg, level) =>
-              val logLevel = LogLevel(level)
-              Logger(pkg).setLogLevel(logLevel)
+              additionalLogLevels += pkg -> LogLevel(level)
             case _ =>
               warn(s"Ignoring invalid argument: ${arg}. Use -L(package)=(log level) to set log levels")
           }
@@ -82,10 +82,22 @@ private[airspec] object AirSpecSbtRunner extends LogSupport {
       i += 1
     }
 
-    new AirSpecSbtRunner(AirSpecConfig(remaining.result(), logLevel), remoteArgs, testClassLoader)
+    new AirSpecSbtRunner(
+      AirSpecConfig(
+        remaining.result(),
+        defaultLogLevel = logLevel,
+        additionalLogLevels = additionalLogLevels.result()
+      ),
+      remoteArgs,
+      testClassLoader
+    )
   }
 
-  case class AirSpecConfig(args: Array[String], defaultLogLevel: LogLevel = LogLevel.INFO) {
+  case class AirSpecConfig(
+      args: Array[String],
+      defaultLogLevel: LogLevel = LogLevel.INFO,
+      additionalLogLevels: Map[String, LogLevel] = Map.empty
+  ) {
     val specMatcher: AirSpecMatcher = {
       // For now, we only support regex-based test name matcher using the first argument
       args.find(x => !x.startsWith("-")) match {
diff --git a/airspec/src/main/scala/wvlet/airspec/runner/AirSpecTaskRunner.scala b/airspec/src/main/scala/wvlet/airspec/runner/AirSpecTaskRunner.scala
index b59e00abcf..c10964332e 100644
--- a/airspec/src/main/scala/wvlet/airspec/runner/AirSpecTaskRunner.scala
+++ b/airspec/src/main/scala/wvlet/airspec/runner/AirSpecTaskRunner.scala
@@ -17,7 +17,7 @@ import sbt.testing.*
 import wvlet.airframe.{Design, Session}
 import wvlet.airspec.runner.AirSpecSbtRunner.AirSpecConfig
 import wvlet.airspec.spi.{AirSpecContext, AirSpecException}
-import wvlet.log.{Logger, LogSupport}
+import wvlet.log.{Logger, LogSupport, LogLevel}
 
 import scala.concurrent.{ExecutionContext, Future, Promise}
 import scala.util.{Failure, Success, Try}
@@ -88,11 +88,19 @@ private[airspec] class AirSpecTaskRunner(
     */
   def runTask: Future[Unit] = {
     val startTimeNanos = System.nanoTime()
-    var prevLogLevel   = Logger(testClassName).getLogLevel
+    val prevLogLevel   = Map.newBuilder[String, LogLevel]
+    prevLogLevel += testClassName -> Logger(testClassName).getLogLevel
     Future
       .apply {
         // Set the default log level for the class
         Logger(testClassName).setLogLevel(config.defaultLogLevel)
+
+        // Set log level for other classes
+        config.additionalLogLevels.foreach { case (pkg, level) =>
+          prevLogLevel += pkg -> Logger(pkg).getLogLevel
+          Logger(pkg).setLogLevel(level)
+        }
+
         // Start a background log level scanner thread. If a thread is already running, reuse it.
         compat.startLogScanner
       }
@@ -110,7 +118,11 @@ private[airspec] class AirSpecTaskRunner(
       }
       .transform { case ret =>
         compat.stopLogScanner
-        Logger(testClassName).setLogLevel(prevLogLevel)
+
+        // Reset log levels
+        prevLogLevel.result().foreach { case (pkg, level) =>
+          Logger(pkg).setLogLevel(level)
+        }
         ret
       }
       .recover { case e: Throwable =>