Rename to CSVFormatter and move format selection to yml

sdv-dev · May 15, 2024 · b552ac6 · b552ac6
1 parent 4355f20
commit b552ac6
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 6 deletions.
diff --git a/sdv/logging/logger.py b/sdv/logging/logger.py
@@ -7,7 +7,7 @@
 from sdv.logging.utils import get_sdv_logger_config
 
 
-class CSVHandler(logging.Formatter):
+class CSVFormatter(logging.Formatter):
     """Logging formatter to convert to CSV."""
 
     def __init__(self):
@@ -59,11 +59,14 @@ def get_sdv_logger(logger_name):
             logger.removeHandler(handler)
 
         if logger_name in logger_conf.get('loggers'):
-            formatter = CSVHandler()
+            formatter = None
             config = logger_conf.get('loggers').get(logger_name)
             log_level = getattr(logging, config.get('level', 'INFO'))
             if config.get('format'):
-                formatter = logging.Formatter(config.get('format'))
+                if config.get('format') == 'CSV':
+                    formatter = CSVFormatter()
+                else:
+                    formatter = logging.Formatter(config.get('format'))
 
             logger.setLevel(log_level)
             logger.propagate = config.get('propagate', False)

diff --git a/sdv/logging/sdv_logger_config.yml b/sdv/logging/sdv_logger_config.yml
@@ -4,24 +4,28 @@ loggers:
   SingleTableSynthesizer:
     level: INFO
     propagate: false
+    format: CSV
     handlers:
       class: logging.FileHandler
       filename: sdv_logs.csv
   MultiTableSynthesizer:
     level: INFO
     propagate: false
+    format: CSV
     handlers:
       class: logging.FileHandler
       filename: sdv_logs.csv
     MultiTableMetadata:
       level: INFO
       propagate: false
+      format: CSV
       handlers:
         class: logging.FileHandler
         filename: sdv_logs.csv
     SingleTableMetadata:
       level: INFO
       propagate: false
+      format: CSV
       handlers:
         class: logging.FileHandler
         filename: sdv_logs.csv
diff --git a/tests/unit/logging/test_logger.py b/tests/unit/logging/test_logger.py
@@ -2,15 +2,15 @@
 import logging
 from unittest.mock import Mock, patch
 
-from sdv.logging.logger import CSVHandler, get_sdv_logger
+from sdv.logging.logger import CSVFormatter, get_sdv_logger
 
 
-class TestCSVHandler:
+class TestCSVFormatter:
 
     def test_format(self):
         """Test CSV formatter correctly formats the log entry."""
         # Setup
-        instance = CSVHandler()
+        instance = CSVFormatter()
         instance.writer = Mock()
         instance.output = Mock()
         record = Mock()
@@ -56,3 +56,41 @@ def test_get_sdv_logger(mock_get_sdv_logger_config, mock_getlogger, mock_streamh
     # Assert
     mock_logger_instance.setLevel.assert_called_once_with(logging.DEBUG)
     mock_logger_instance.addHandler.assert_called_once()
+
+
+@patch('sdv.logging.logger.CSVFormatter')
+@patch('sdv.logging.logger.logging.FileHandler')
+@patch('sdv.logging.logger.logging.getLogger')
+@patch('sdv.logging.logger.get_sdv_logger_config')
+def test_get_sdv_logger_csv(mock_get_sdv_logger_config, mock_getlogger,
+                            mock_filehandler, mock_csvformatter):
+    # Setup
+    mock_logger_conf = {
+        'log_registry': 'local',
+        'loggers': {
+            'test_logger': {
+                'level': 'DEBUG',
+                'format': 'CSV',
+                'handlers': {
+                    'filename': 'logfile.csv',
+                    'class': 'logging.FileHandler'
+                }
+            }
+        }
+    }
+    mock_get_sdv_logger_config.return_value = mock_logger_conf
+    mock_logger_instance = Mock()
+    mock_logger_instance.handlers = []
+    mock_getlogger.return_value = mock_logger_instance
+    mock_filehandler_instance = Mock()
+    mock_filehandler.return_value = mock_filehandler_instance
+
+    # Run
+    get_sdv_logger('test_logger')
+
+    # Assert
+    mock_logger_instance.setLevel.assert_called_once_with(logging.DEBUG)
+    mock_logger_instance.addHandler.assert_called_once_with(mock_filehandler_instance)
+    mock_filehandler.assert_called_once_with('logfile.csv')
+    mock_filehandler_instance.setLevel.assert_called_once_with(logging.DEBUG)
+    mock_filehandler_instance.setFormatter.assert_called_once_with(mock_csvformatter.return_value)