Merge pull request #26089 from edx/robrap/fix-cookie-monitoring-middleware

fix cookie monitoring middleware

openedx/core/lib/request_utils.py

@@ -153,54 +153,61 @@ class CookieMonitoringMiddleware(MiddlewareMixin):
                 # Add or update the size for this group.
                 cookie_groups_to_size[grouping_name] = cookie_groups_to_size.get(grouping_name, 0) + cookie_size
 
-        max_cookie_name = max(cookie_names_to_size, key=lambda name: cookie_names_to_size[name])
-        max_cookie_size = cookie_names_to_size[max_cookie_name]
+        if cookie_names_to_size:
+            self.set_custom_attributes_for_top_n(
+                cookie_names_to_size,
+                top_n_cookies_captured,
+                attribute_prefix='cookies',
+            )
+
+            max_cookie_name = max(cookie_names_to_size, key=lambda name: cookie_names_to_size[name])
+            max_cookie_size = cookie_names_to_size[max_cookie_name]
+
+            set_custom_attribute('cookies.max.name', max_cookie_name)
+            set_custom_attribute('cookies.max.size', max_cookie_size)
+
+        if cookie_groups_to_size:
+            self.set_custom_attributes_for_top_n(
+                cookie_groups_to_size,
+                top_n_cookie_groups_captured,
+                attribute_prefix='cookies.group',
+            )
+
+            max_group_cookie_name = max(cookie_groups_to_size, key=lambda name: cookie_groups_to_size[name])
+            max_group_cookie_size = cookie_groups_to_size[max_group_cookie_name]
+
+            # If a single cookies is bigger than any group of cookies, we want max_group... to reflect that.
+            # Treating an individual cookie as a group of 1 for calculating the max.
+            if max_group_cookie_size < max_cookie_size:
+                max_group_cookie_name = max_cookie_name
+                max_group_cookie_size = max_cookie_size
+
+            set_custom_attribute('cookies.max.group.name', max_group_cookie_name)
+            set_custom_attribute('cookies.max.group.size', max_group_cookie_size)
 
-        max_group_cookie_name = max(cookie_groups_to_size, key=lambda name: cookie_groups_to_size[name])
-        max_group_cookie_size = cookie_groups_to_size[max_group_cookie_name]
+        total_cookie_size = sum(cookie_names_to_size.values())
+        set_custom_attribute('cookies_total_size', total_cookie_size)
+        log.debug(u'cookies_total_size = %d', total_cookie_size)
 
-        # If a single cookies is bigger than any group of cookies, we want max_group... to reflect that.
-        # Treating an individual cookie as a group of 1 for calculating the max.
-        if max_group_cookie_size < max_cookie_size:
-            max_group_cookie_name = max_cookie_name
-            max_group_cookie_size = max_cookie_size
+    def set_custom_attributes_for_top_n(self, names_to_size, top_n_captured, attribute_prefix):
+        """
+        Sets custom metric for the top N biggest cookies or cookie groups.
 
-        # Log only the top N biggest cookies.
+        Arguments:
+            names_to_size: Dict of sizes keyed by cookie name or cookie group name
+            top_n_captured: Number of largest sizes to monitor.
+            attribute_prefix: Prefix (cookies|cookies.group) to use in the custom attribute name.
+        """
         top_n_cookies = sorted(
-            cookie_names_to_size,
-            key=lambda x: cookie_names_to_size[x],
+            names_to_size,
+            key=lambda x: names_to_size[x],
             reverse=True,
-        )[:top_n_cookies_captured]
+        )[:top_n_captured]
         for index, name in enumerate(top_n_cookies, start=1):
-            size = cookie_names_to_size[name]
-            name_attribute = 'cookies.{}.name'.format(index)
-            size_attribute = 'cookies.{}.size'.format(index)
-
-            set_custom_attribute(name_attribute, name)
-            set_custom_attribute(size_attribute, size)
-            log.debug(u'%s = %d', name, size)
-
-        # Log only the top N biggest groups.
-        top_n_cookie_groups = sorted(
-            cookie_groups_to_size,
-            key=lambda x: cookie_groups_to_size[x],
-            reverse=True,
-        )[:top_n_cookie_groups_captured]
-
-        for index, name in enumerate(top_n_cookie_groups, start=1):
-            size = cookie_groups_to_size[name]
-            name_attribute = 'cookies.group.{}.name'.format(index)
-            size_attribute = 'cookies.group.{}.size'.format(index)
+            size = names_to_size[name]
+            name_attribute = '{}.{}.name'.format(attribute_prefix, index)
+            size_attribute = '{}.{}.size'.format(attribute_prefix, index)
 
             set_custom_attribute(name_attribute, name)
             set_custom_attribute(size_attribute, size)
-            log.debug(u'%s = %d', name, size)
-
-        set_custom_attribute('cookies.max.name', max_cookie_name)
-        set_custom_attribute('cookies.max.size', max_cookie_size)
-        set_custom_attribute('cookies.max.group.name', max_group_cookie_name)
-        set_custom_attribute('cookies.max.group.size', max_group_cookie_size)
-
-        total_cookie_size = sum(cookie_names_to_size.values())
-        set_custom_attribute('cookies_total_size', total_cookie_size)
-        log.debug(u'cookies_total_size = %d', total_cookie_size)
+            log.debug('%s = %d', name, size)

openedx/core/lib/tests/test_request_utils.py

@@ -131,7 +131,7 @@ class RequestUtilTestCase(unittest.TestCase):
             call('cookies.max.group.name', 'a'),
             call('cookies.max.group.size', 100),
             call('cookies_total_size', 189),
-        ])
+        ], any_order=True)
 
     @patch("openedx.core.lib.request_utils.CAPTURE_COOKIE_SIZES")
     @patch("openedx.core.lib.request_utils.set_custom_attribute")
@@ -163,4 +163,43 @@ class RequestUtilTestCase(unittest.TestCase):
             call('cookies.max.group.name', 'b'),
             call('cookies.max.group.size', 35),
             call('cookies_total_size', 45)
-        ])
+        ], any_order=True)
+
+    @patch("openedx.core.lib.request_utils.CAPTURE_COOKIE_SIZES")
+    @patch("openedx.core.lib.request_utils.set_custom_attribute")
+    def test_cookie_monitoring_no_cookies(self, mock_set_custom_attribute, mock_capture_cookie_sizes):
+
+        mock_capture_cookie_sizes.is_enabled.return_value = True
+        middleware = CookieMonitoringMiddleware()
+
+        mock_request = Mock()
+        mock_request.COOKIES = {}
+
+        middleware.process_request(mock_request)
+
+        mock_set_custom_attribute.assert_has_calls([call('cookies_total_size', 0)], any_order=True)
+
+    @patch("openedx.core.lib.request_utils.CAPTURE_COOKIE_SIZES")
+    @patch("openedx.core.lib.request_utils.set_custom_attribute")
+    def test_cookie_monitoring_no_groups(self, mock_set_custom_attribute, mock_capture_cookie_sizes):
+
+        mock_capture_cookie_sizes.is_enabled.return_value = True
+        middleware = CookieMonitoringMiddleware()
+
+        mock_request = Mock()
+        mock_request.COOKIES = {
+            "a": "." * 10,
+            "b": "." * 15,
+        }
+
+        middleware.process_request(mock_request)
+
+        mock_set_custom_attribute.assert_has_calls([
+            call('cookies.max.name', 'b'),
+            call('cookies.max.size', 15),
+            call('cookies.1.name', 'b'),
+            call('cookies.1.size', 15),
+            call('cookies.2.name', 'a'),
+            call('cookies.2.size', 10),
+            call('cookies_total_size', 25),
+        ], any_order=True)