Просмотр исходного кода

Fixes #22923: Fix post-exception cleanup under event_tracking() (#22926)

Jeremy Stretch 1 неделя назад
Родитель
Сommit
93f16a536d
2 измененных файлов с 101 добавлено и 11 удалено
  1. 14 11
      netbox/netbox/context_managers.py
  2. 87 0
      netbox/netbox/tests/test_context_managers.py

+ 14 - 11
netbox/netbox/context_managers.py

@@ -15,17 +15,20 @@ def event_tracking(request):
 
     :param request: WSGIRequest object with a unique `id` set
     """
-    current_request.set(request)
-    events_queue.set({})
-    query_cache.set(defaultdict(dict))
+    request_token = current_request.set(request)
+    queue_token = events_queue.set({})
+    cache_token = query_cache.set(defaultdict(dict))
 
-    yield
+    try:
+        yield
 
-    # Flush queued webhooks to RQ
-    if events := list(events_queue.get().values()):
-        flush_events(events)
+        # Flush queued webhooks to RQ. This is done only if the wrapped block completed successfully; events
+        # queued by a failed request or job must not be dispatched.
+        if events := list(events_queue.get().values()):
+            flush_events(events)
 
-    # Clear context vars
-    current_request.set(None)
-    events_queue.set({})
-    query_cache.set(None)
+    finally:
+        # Restore the previous context vars, whether or not the wrapped block raised an exception
+        current_request.reset(request_token)
+        events_queue.reset(queue_token)
+        query_cache.reset(cache_token)

+ 87 - 0
netbox/netbox/tests/test_context_managers.py

@@ -0,0 +1,87 @@
+import uuid
+from unittest.mock import patch
+
+from django.test import RequestFactory, SimpleTestCase
+
+from netbox.context import current_request, events_queue, query_cache
+from netbox.context_managers import event_tracking
+
+
+def _build_request():
+    request = RequestFactory().get('/')
+    request.id = uuid.uuid4()
+    request.user = None
+    return request
+
+
+class EventTrackingTestCase(SimpleTestCase):
+    """
+    Verify that event_tracking() populates and restores its context variables.
+    """
+    def assertContextVarsRestored(self):
+        self.assertIsNone(current_request.get())
+        self.assertEqual(events_queue.get(), {})
+        self.assertIsNone(query_cache.get())
+
+    def test_context_vars_set_within_block(self):
+        request = _build_request()
+
+        with event_tracking(request):
+            self.assertIs(current_request.get(), request)
+            self.assertEqual(events_queue.get(), {})
+            self.assertIsNotNone(query_cache.get())
+
+        self.assertContextVarsRestored()
+
+    def test_context_vars_restored_after_exception(self):
+        request = _build_request()
+
+        with self.assertRaises(RuntimeError):
+            with event_tracking(request):
+                raise RuntimeError('simulated view/script failure')
+
+        self.assertContextVarsRestored()
+
+    def test_events_flushed_on_success(self):
+        request = _build_request()
+
+        with patch('netbox.context_managers.flush_events') as flush_events:
+            with event_tracking(request):
+                events_queue.get()['foo'] = 'bar'
+
+        flush_events.assert_called_once_with(['bar'])
+
+    def test_events_not_flushed_after_exception(self):
+        request = _build_request()
+
+        with patch('netbox.context_managers.flush_events') as flush_events:
+            with self.assertRaises(RuntimeError):
+                with event_tracking(request):
+                    events_queue.get()['foo'] = 'bar'
+                    raise RuntimeError('simulated view/script failure')
+
+        flush_events.assert_not_called()
+
+    def test_nested_context_restores_outer_values(self):
+        outer_request = _build_request()
+        inner_request = _build_request()
+
+        with patch('netbox.context_managers.flush_events'):
+            with event_tracking(outer_request):
+                outer_cache = query_cache.get()
+                outer_queue = events_queue.get()
+                outer_queue['outer'] = 'event'
+
+                with event_tracking(inner_request):
+                    self.assertIs(current_request.get(), inner_request)
+                    self.assertIsNot(events_queue.get(), outer_queue)
+                    self.assertEqual(events_queue.get(), {})
+
+                # The outer request's context must be restored intact, including any events it had
+                # already queued
+                self.assertIs(current_request.get(), outer_request)
+                self.assertIs(query_cache.get(), outer_cache)
+                self.assertIs(events_queue.get(), outer_queue)
+                self.assertEqual(events_queue.get(), {'outer': 'event'})
+
+        self.assertContextVarsRestored()