From 027f0862fd75294ce85c4eb5543845f7409f2aab Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Sun, 2 Mar 2025 18:49:18 -0800
Subject: [PATCH] fix test

---
 llama_stack/logcat.py |  5 ++++-
 tests/test_logcat.py  | 22 +++++++---------------
 2 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/llama_stack/logcat.py b/llama_stack/logcat.py
index b9314ec2e..0e11cb782 100644
--- a/llama_stack/logcat.py
+++ b/llama_stack/logcat.py
@@ -161,7 +161,10 @@ def init(default_level: int = logging.INFO) -> None:
 
 
 def _should_log(level: int, category: str) -> bool:
-    category_level = _category_levels.get(category.lower(), _default_level)
+    category = category.lower()
+    if category not in _category_levels:
+        return False
+    category_level = _category_levels[category]
     return level >= category_level
 
 
diff --git a/tests/test_logcat.py b/tests/test_logcat.py
index bcdcb33ab..4a116a08f 100644
--- a/tests/test_logcat.py
+++ b/tests/test_logcat.py
@@ -16,15 +16,8 @@ class TestLogcat(unittest.TestCase):
     def setUp(self):
         self.original_env = os.environ.get("LLAMA_STACK_LOGGING")
 
-        self.original_handlers = list(logcat._logger.handlers)
-        self.original_level = logcat._logger.level
-
         self.log_output = io.StringIO()
-        self.handler = logging.StreamHandler(self.log_output)
-        self.handler.setFormatter(logging.Formatter("[%(category)s] %(message)s"))
-
-        logcat._logger.handlers.clear()
-        logcat._logger.addHandler(self.handler)
+        self._init_logcat()
 
     def tearDown(self):
         if self.original_env is not None:
@@ -32,12 +25,12 @@ class TestLogcat(unittest.TestCase):
         else:
             os.environ.pop("LLAMA_STACK_LOGGING", None)
 
+    def _init_logcat(self):
+        logcat.init(default_level=logging.DEBUG)
+        self.handler = logging.StreamHandler(self.log_output)
+        self.handler.setFormatter(logging.Formatter("[%(category)s] %(message)s"))
         logcat._logger.handlers.clear()
-        for handler in self.original_handlers:
-            logcat._logger.addHandler(handler)
-        logcat._logger.setLevel(self.original_level)
-
-        logcat._initialize()
+        logcat._logger.addHandler(self.handler)
 
     def test_basic_logging(self):
         logcat.info("server", "Info message")
@@ -62,8 +55,7 @@ class TestLogcat(unittest.TestCase):
 
     def test_env_var_control(self):
         os.environ["LLAMA_STACK_LOGGING"] = "server=debug;inference=warning"
-
-        logcat._initialize()
+        self._init_logcat()
 
         # These should be visible based on the environment settings
         logcat.debug("server", "Server debug message")