refactor: make IsolateLogger more explicit (#152)

* refactor: make IsolateLogger more explicit

* Update src/isolate/logger.py

Co-authored-by: Ruslan Kuprieiev <[email protected]>

* typing

* tests

---------

Co-authored-by: Ruslan Kuprieiev <[email protected]>

src/isolate/logger.py

@@ -1,24 +1,18 @@
 import json
 import os
+from typing import Dict
+
+from isolate.logs import LogLevel, LogSource
 
 
 # NOTE: we probably should've created a proper `logging.getLogger` here,
 # but it handling `source` would be not trivial, so we are better off
 # just keeping it simple for now.
 class IsolateLogger:
-    def __init__(self):
-        self.log_labels = {}
-        raw = os.getenv("ISOLATE_LOG_LABELS")
-        if raw:
-            labels = json.loads(raw)
-            for key, value in labels.items():
-                if value.startswith("$"):
-                    expanded = os.getenv(value[1:])
-                else:
-                    expanded = value
-                self.log_labels[key] = expanded
-
-    def log(self, level, message, source):
+    def __init__(self, log_labels: Dict[str, str]):
+        self.log_labels = log_labels
+
+    def log(self, level: LogLevel, message: str, source: LogSource) -> None:
         record = {
             "isolate_source": source.name,
             "level": level.name,
@@ -27,5 +21,29 @@ def log(self, level, message, source):
         }
         print(json.dumps(record))
 
+    @classmethod
+    def with_env_expanded(cls, labels: Dict[str, str]) -> "IsolateLogger":
+        for key, value in labels.items():
+            if value.startswith("$"):
+                expanded = os.getenv(value[1:])
+            else:
+                expanded = value
+            if expanded is not None:
+                labels[key] = expanded
+
+        return cls(labels)
+
+    @classmethod
+    def from_env(cls) -> "IsolateLogger":
+        _labels: Dict[str, str] = {}
+        raw = os.getenv("ISOLATE_LOG_LABELS")
+        if raw:
+            try:
+                _labels = json.loads(raw)
+            except json.JSONDecodeError:
+                print("Failed to parse ISOLATE_LOG_LABELS")
+
+        return cls.with_env_expanded(labels=_labels)
+
 
-logger = IsolateLogger()
+ENV_LOGGER = IsolateLogger.from_env()

src/isolate/server/server.py

@@ -29,7 +29,7 @@
 from isolate.backends.virtualenv import VirtualPythonEnvironment
 from isolate.connections.grpc import AgentError, LocalPythonGRPC
 from isolate.connections.grpc.configuration import get_default_options
-from isolate.logger import logger
+from isolate.logger import ENV_LOGGER, IsolateLogger
 from isolate.logs import Log, LogLevel, LogSource
 from isolate.server import definitions, health
 from isolate.server.health_server import HealthServicer
@@ -453,9 +453,10 @@ def _proxy_to_queue(
 @dataclass
 class LogHandler:
     messages: Queue
+    logger: IsolateLogger = ENV_LOGGER
 
     def handle(self, log: Log) -> None:
-        logger.log(log.level, log.message, source=log.source)
+        self.logger.log(log.level, log.message, source=log.source)
         self._add_log_to_queue(log)
 
     def _add_log_to_queue(self, log: Log) -> None:

tests/test_logger.py

@@ -1,22 +1,45 @@
 import json
 
+import pytest
 from isolate.logger import IsolateLogger
 
 
-def test_logger(monkeypatch):
-    labels = {
+@pytest.fixture
+def log_labels():
+    return {
         "foo": "$MYENVVAR1",
         "bar": "$MYENVVAR2",
         "baz": "baz",
-        "qux": "qux",
+        "qux": "$NOTTHERE",
     }
+
+
+def test_logger_with_env_expanded(log_labels, monkeypatch):
     monkeypatch.setenv("MYENVVAR1", "myenvvar1")
     monkeypatch.setenv("MYENVVAR2", "myenvvar2")
-    monkeypatch.setenv("ISOLATE_LOG_LABELS", json.dumps(labels))
-    logger = IsolateLogger()
+    logger = IsolateLogger.with_env_expanded(log_labels)
     assert logger.log_labels == {
         "foo": "myenvvar1",
         "bar": "myenvvar2",
         "baz": "baz",
-        "qux": "qux",
+        "qux": "$NOTTHERE",
     }
+
+
+def test_logger_from_env(log_labels, monkeypatch):
+    monkeypatch.setenv("MYENVVAR1", "myenvvar1")
+    monkeypatch.setenv("MYENVVAR2", "myenvvar2")
+    monkeypatch.setenv("ISOLATE_LOG_LABELS", json.dumps(log_labels))
+    logger = IsolateLogger.from_env()
+    assert logger.log_labels == {
+        "foo": "myenvvar1",
+        "bar": "myenvvar2",
+        "baz": "baz",
+        "qux": "$NOTTHERE",
+    }
+
+
+def test_logger_direct(log_labels):
+    logger = IsolateLogger(log_labels=log_labels)
+    # should not do env expansion
+    assert logger.log_labels == log_labels