From 6941b10ba5d59200f939c2126ad6e227e247c43e Mon Sep 17 00:00:00 2001
From: Viktor Taskov <viktor.taskov@seldon.io>
Date: Fri, 21 Jul 2023 13:22:28 +0100
Subject: [PATCH] Verify structured logging format for model details

---
 mlserver/logging.py   | 10 +++++-----
 tests/test_logging.py | 37 +++++++++++++++++++++----------------
 2 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/mlserver/logging.py b/mlserver/logging.py
index 98694074a..a02af6178 100644
--- a/mlserver/logging.py
+++ b/mlserver/logging.py
@@ -42,7 +42,9 @@ class ModelLoggerFormatter(logging.Formatter):
     """
 
     _UNSTRUCTURED_FORMAT = "%(asctime)s [%(name)s]%(model)s %(levelname)s - %(message)s"
-    _STRUCTURED_FORMAT = "{'time':'%(asctime)s', 'name': '%(name)s', 'level': '%(levelname)s', 'message': '%(message)s' %(model)s}"
+    _STRUCTURED_FORMAT = "{\"time\":\"%(asctime)s\", \"name\": \"%(name)s\", " \
+                         "\"level\": \"%(levelname)s\", \"message\": \"%(message)s\"" \
+                         " %(model)s}"
 
     def __init__(self, use_structured_logging: bool):
         super().__init__(
@@ -65,11 +67,9 @@ def _format_unstructured_model_details(name: str, version: str) -> str:
     def _format_structured_model_details(name: str, version: str) -> str:
         if not name:
             return ""
-        model_details = []
-        model_details.append(f"'model_name': {name}")
-        print()
+        model_details = f", \"model_name\": \"{name}\""
         if version:
-            model_details += f',"model_version": "{version}"'
+            model_details += f", \"model_version\": \"{version}\""
         return model_details
 
     def format(self, record: logging.LogRecord) -> str:
diff --git a/tests/test_logging.py b/tests/test_logging.py
index 94b065f6b..6b04ffb08 100644
--- a/tests/test_logging.py
+++ b/tests/test_logging.py
@@ -48,7 +48,7 @@ def test_model_logging_formatter_unstructured(name, version, expected_fmt, caplo
     logger.info("After model context")
 
     log_records = caplog.text.strip().split("\n")
-    assert len(caplog.messages) == 3
+    assert len(log_records) == 3
 
     assert expected_fmt not in log_records[0]
     assert expected_fmt in log_records[1]
@@ -56,26 +56,26 @@ def test_model_logging_formatter_unstructured(name, version, expected_fmt, caplo
 
 
 @pytest.mark.parametrize(
-    "name, version, expected_fmt",
+    "name, version, expected_model_details",
     [
         (
             "foo",
             "v1.0",
-            "{}",
+            {"model_name", "model_version"}
+        ),
+        (
+            "foo",
+            "",
+            {"model_name"}
+        ),
+        (
+            "foo",
+            None,
+            {"model_name"}
         ),
-        # (
-        #     "foo",
-        #     "",
-        #     "[foo]",
-        # ),
-        # (
-        #     "foo",
-        #     None,
-        #     "[foo]",
-        # ),
     ],
 )
-def test_model_logging_formatter_structured(name, version, expected_fmt, caplog):
+def test_model_logging_formatter_structured(name, version, expected_model_details, caplog):
     caplog.handler.setFormatter(ModelLoggerFormatter(use_structured_logging=True))
     caplog.set_level(INFO)
 
@@ -88,9 +88,14 @@ def test_model_logging_formatter_structured(name, version, expected_fmt, caplog)
         logger.info("Inside model context")
     logger.info("After model context")
 
-    structured_log = json.loads(caplog.text)
+    log_records = caplog.text.strip().split("\n")
+    assert len(log_records) == 3
+
+    json_log_records = [json.loads(lr) for lr in log_records]
 
-    assert "model_name" in structured_log
+    assert not expected_model_details <= json_log_records[0].keys()
+    assert expected_model_details <= json_log_records[1].keys()
+    assert not expected_model_details <= json_log_records[2].keys()
 
 
 @pytest.mark.parametrize("debug", [True, False])