Expose the method to set new log prefix in PrefixLoggerAdapter (#852)

2 files changed, 26 insertions, 3 deletions

diff --git a/mobly/logger.py b/mobly/logger.py
index 365f5c5..ea1de12 100644
--- a/mobly/logger.py
+++ b/mobly/logger.py
@@ -412,11 +412,20 @@ class PrefixLoggerAdapter(logging.LoggerAdapter):
     """Processes the logging call to insert contextual information.
 
     Args:
-      msg: the logging message
-      kwargs: keyword arguments passed in to a logging call
+      msg: The logging message.
+      kwargs: Keyword arguments passed in to a logging call.
 
     Returns:
-      the message and kwargs modified.
+      The message and kwargs modified.
     """
     new_msg = f'{self.extra[PrefixLoggerAdapter.EXTRA_KEY_LOG_PREFIX]} {msg}'
     return (new_msg, kwargs)
+
+  def set_log_prefix(self, prefix: str) -> None:
+    """Sets the log prefix to the given string.
+
+    Args:
+      prefix: The new log prefix.
+    """
+    self.debug('Setting the log prefix to "%s".', prefix)
+    self.extra[PrefixLoggerAdapter.EXTRA_KEY_LOG_PREFIX] = prefix
diff --git a/tests/mobly/logger_test.py b/tests/mobly/logger_test.py
index 8a1e778..e0ac14d 100755
--- a/tests/mobly/logger_test.py
+++ b/tests/mobly/logger_test.py
@@ -223,6 +223,20 @@ class LoggerTest(unittest.TestCase):
     self.assertEqual(processed_log, '[MOCK_PREFIX] mock log line')
     self.assertIs(processed_kwargs, kwargs)
 
+  def test_prefix_logger_adapter_modify_prefix(self):
+    extra = {
+        logger.PrefixLoggerAdapter.EXTRA_KEY_LOG_PREFIX: 'MOCK_PREFIX',
+    }
+    adapted_logger = logger.PrefixLoggerAdapter(mock.Mock(), extra)
+    adapted_logger.set_log_prefix('[NEW]')
+
+    kwargs = mock.Mock()
+    processed_log, processed_kwargs = adapted_logger.process('mock log line',
+                                                             kwargs=kwargs)
+
+    self.assertEqual(processed_log, '[NEW] mock log line')
+    self.assertIs(processed_kwargs, kwargs)
+
 
 if __name__ == "__main__":
   unittest.main()