From 0a06fdbfb58189f69faadaf4ccbb4c176bf9314a Mon Sep 17 00:00:00 2001
From: "stone-24tch3r (aider)"
 <100294019+stone-w4tch3r@users.noreply.github.com>
Date: Sun, 22 Sep 2024 22:37:25 +0500
Subject: [PATCH] feat: Add tests for trace_id_handler

---
 tests/test_trace_id_handler.py | 41 ++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 tests/test_trace_id_handler.py

diff --git a/tests/test_trace_id_handler.py b/tests/test_trace_id_handler.py
new file mode 100644
index 0000000..df0e961
--- /dev/null
+++ b/tests/test_trace_id_handler.py
@@ -0,0 +1,41 @@
+import unittest
+import uuid
+from unittest.mock import patch
+from core.trace_id_handler import TraceIdHandler
+
+class TestTraceIdHandler(unittest.TestCase):
+
+    def test_logging_scope(self):
+        with TraceIdHandler.logging_scope():
+            trace_id = TraceIdHandler.get_current_trace_id()
+            self.assertIsInstance(trace_id, uuid.UUID)
+
+        # After exiting the scope, trace_id should be None
+        self.assertIsNone(TraceIdHandler.get_current_trace_id())
+
+    def test_nested_logging_scopes(self):
+        with TraceIdHandler.logging_scope() as outer_trace_id:
+            with TraceIdHandler.logging_scope() as inner_trace_id:
+                self.assertNotEqual(outer_trace_id, inner_trace_id)
+                self.assertEqual(TraceIdHandler.get_current_trace_id(), inner_trace_id)
+            
+            # After exiting inner scope, we should be back to outer trace_id
+            self.assertEqual(TraceIdHandler.get_current_trace_id(), outer_trace_id)
+
+        # After exiting all scopes, trace_id should be None
+        self.assertIsNone(TraceIdHandler.get_current_trace_id())
+
+    @patch('uuid.uuid4')
+    def test_consistent_trace_id_within_scope(self, mock_uuid4):
+        mock_uuid = uuid.UUID('12345678-1234-5678-1234-567812345678')
+        mock_uuid4.return_value = mock_uuid
+
+        with TraceIdHandler.logging_scope():
+            self.assertEqual(TraceIdHandler.get_current_trace_id(), mock_uuid)
+            self.assertEqual(TraceIdHandler.get_current_trace_id(), mock_uuid)  # Should return the same UUID
+
+    def test_get_current_trace_id_outside_scope(self):
+        self.assertIsNone(TraceIdHandler.get_current_trace_id())
+
+if __name__ == '__main__':
+    unittest.main()