[MINOR][TESTS] Use SparkContext.setLogLevel in QuietTest

### What changes were proposed in this pull request? This PR proposes to use SparkContext.setLogLevel in QuietTest ### Why are the changes needed? To avoid the direct update of logger but reuses SparkContext.setLogLevel. ### Does this PR introduce _any_ user-facing change? No, test-only. ### How was this patch tested? Manually tested. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #48966 from HyukjinKwon/minor-log-level. Authored-by: Hyukjin Kwon <[email protected]> Signed-off-by: Hyukjin Kwon <[email protected]>
apache · Nov 26, 2024 · 9bab16a · 9bab16a
1 parent 9054aa2
commit 9bab16a
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/python/pyspark/testing/utils.py b/python/pyspark/testing/utils.py
@@ -192,14 +192,16 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
 
 class QuietTest:
     def __init__(self, sc):
-        self.log4j = sc._jvm.org.apache.log4j
+        self.sc = sc
 
     def __enter__(self):
-        self.old_level = self.log4j.LogManager.getRootLogger().getLevel()
-        self.log4j.LogManager.getRootLogger().setLevel(self.log4j.Level.FATAL)
+        self.old_level = (
+            self.sc._jvm.org.apache.log4j.LogManager.getRootLogger().getLevel().toString()
+        )
+        self.sc.setLogLevel("FATAL")
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.log4j.LogManager.getRootLogger().setLevel(self.old_level)
+        self.sc.setLogLevel(self.old_level)
 
 
 class PySparkTestCase(unittest.TestCase):