From f94628db61ae84c6ef90d706e88afeac8b9889b9 Mon Sep 17 00:00:00 2001
From: Eduard Karacharov <eduard.karacharov@semrush.com>
Date: Mon, 13 May 2024 11:11:09 +0300
Subject: [PATCH] log message format test coverage

---
 tests/test_blocks.py | 22 ++++++++++++++++++++++
 tests/util.py        | 15 +++++++++++----
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/tests/test_blocks.py b/tests/test_blocks.py
index ff992678..6129d427 100644
--- a/tests/test_blocks.py
+++ b/tests/test_blocks.py
@@ -1,3 +1,4 @@
+import re
 import types
 from unittest.mock import patch
 
@@ -245,3 +246,24 @@ def test_logs_with_compression(self):
                 query = 'SELECT 1'
                 client.execute(query, settings=settings)
                 self.assertIn(query, buffer.getvalue())
+
+    def test_logs_message_format(self):
+        log_regex_source = (
+            r'^\[ .+? \] '                 # hostname
+            r'[\d-]{10} [\d:]{8}\.\d{6} '  # timestamp
+            r'\[ \d+? \] '                 # thread_id
+            r'\{[\w-]{36}\} '              # query_id
+            r'<[a-z]+?> '                  # priority
+            r'\w+?:'                       # source
+        )
+        log_regex = re.compile(log_regex_source, re.I)
+
+        with capture_logging(
+            'clickhouse_driver.log', 'INFO', as_queue=True
+        ) as queue:
+            settings = {'send_logs_level': 'debug'}
+            query = 'SELECT 1'
+            self.client.execute(query, settings=settings)
+            while not queue.empty():
+                record = queue.get()
+                self.assertRegex(record.message, log_regex)
diff --git a/tests/util.py b/tests/util.py
index db48166c..26d5c858 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -1,8 +1,10 @@
+import logging.handlers
 import os
 from contextlib import contextmanager
 from functools import wraps
 import logging
 from io import StringIO
+from queue import SimpleQueue
 from time import tzset
 from unittest.mock import patch
 
@@ -36,21 +38,26 @@ def wrapper(*args, **kwargs):
 
 
 class LoggingCapturer(object):
-    def __init__(self, logger_name, level):
+    def __init__(self, logger_name, level, as_queue=False):
         self.old_stdout_handlers = []
         self.logger = logging.getLogger(logger_name)
         self.level = level
+        self.as_queue = as_queue
         super(LoggingCapturer, self).__init__()
 
     def __enter__(self):
-        buffer = StringIO()
+        if self.as_queue:
+            store = SimpleQueue()
+            self.new_handler = logging.handlers.QueueHandler(store)
+        else:
+            store = StringIO()
+            self.new_handler = logging.StreamHandler(store)
 
-        self.new_handler = logging.StreamHandler(buffer)
         self.logger.addHandler(self.new_handler)
         self.old_logger_level = self.logger.level
         self.logger.setLevel(self.level)
 
-        return buffer
+        return store
 
     def __exit__(self, *exc_info):
         self.logger.setLevel(self.old_logger_level)