From 7e344e6e0adce4d8e31420363aa08e6e8d3d41fb Mon Sep 17 00:00:00 2001
From: Dan Helfman <witten@torsion.org>
Date: Mon, 27 Jan 2025 12:25:28 -0800
Subject: [PATCH] Complete test coverage for new code (#635).

---
 borgmatic/logger.py       |  2 +-
 tests/unit/test_logger.py | 77 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/borgmatic/logger.py b/borgmatic/logger.py
index 4acaf50d..a81541d1 100644
--- a/borgmatic/logger.py
+++ b/borgmatic/logger.py
@@ -203,7 +203,7 @@ def get_log_prefix():
 def set_log_prefix(prefix):
     '''
     Given a log prefix as a string, set it into the defaults for the formatters on all logging
-    handlers.
+    handlers. Note that this overwrites any existing defaults.
     '''
     for handler in logging.getLogger().handlers:
         try:
diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py
index 5a4d60c9..6d8b0049 100644
--- a/tests/unit/test_logger.py
+++ b/tests/unit/test_logger.py
@@ -276,9 +276,7 @@ def test_get_log_prefix_with_no_prefix_does_not_raise():
         flexmock(
             handlers=[
                 flexmock(
-                    formatter=flexmock(
-                        _style=flexmock(_defaults=flexmock(get=lambda name: None))
-                    )
+                    formatter=flexmock(_style=flexmock(_defaults=flexmock(get=lambda name: None)))
                 ),
                 flexmock(),
             ],
@@ -289,6 +287,79 @@ def test_get_log_prefix_with_no_prefix_does_not_raise():
     assert module.get_log_prefix() == None
 
 
+def test_set_log_prefix_updates_all_handlers():
+    styles = (
+        flexmock(_defaults=None),
+        flexmock(_defaults=None),
+    )
+
+    flexmock(module.logging).should_receive('getLogger').and_return(
+        flexmock(
+            handlers=[
+                flexmock(
+                    formatter=flexmock(
+                        _style=styles[0],
+                    )
+                ),
+                flexmock(
+                    formatter=flexmock(
+                        _style=styles[1],
+                    )
+                ),
+            ],
+            removeHandler=lambda handler: None,
+        )
+    )
+
+    module.set_log_prefix('myprefix')
+
+    for style in styles:
+        assert style._defaults == {'prefix': 'myprefix: '}
+
+
+def test_set_log_prefix_skips_handlers_without_a_formatter():
+    style = flexmock(_defaults=None)
+
+    flexmock(module.logging).should_receive('getLogger').and_return(
+        flexmock(
+            handlers=[
+                flexmock(
+                    formatter=None,
+                ),
+                flexmock(),
+                flexmock(
+                    formatter=flexmock(
+                        _style=style,
+                    )
+                ),
+            ],
+            removeHandler=lambda handler: None,
+        )
+    )
+
+    module.set_log_prefix('myprefix')
+
+    assert style._defaults == {'prefix': 'myprefix: '}
+
+
+def test_log_prefix_sets_prefix_and_then_restores_no_prefix_after():
+    flexmock(module).should_receive('get_log_prefix').and_return(None)
+    flexmock(module).should_receive('set_log_prefix').with_args('myprefix').once()
+    flexmock(module).should_receive('set_log_prefix').with_args(None).once()
+
+    with module.Log_prefix('myprefix'):
+        pass
+
+
+def test_log_prefix_sets_prefix_and_then_restores_original_prefix_after():
+    flexmock(module).should_receive('get_log_prefix').and_return('original')
+    flexmock(module).should_receive('set_log_prefix').with_args('myprefix').once()
+    flexmock(module).should_receive('set_log_prefix').with_args('original').once()
+
+    with module.Log_prefix('myprefix'):
+        pass
+
+
 def test_configure_logging_with_syslog_log_level_probes_for_log_socket_on_linux():
     flexmock(module).should_receive('add_custom_log_levels')
     flexmock(module.logging).ANSWER = module.ANSWER