Custom django-waffle middleware to store testing flags

open_humans/middleware.py

@@ -5,10 +5,16 @@
 from django.conf import settings
 from django.http import HttpResponseRedirect
 
+from waffle import get_waffle_flag_model
+from waffle.middleware import WaffleMiddleware
+from waffle.utils import get_setting as get_waffle_setting
+
 from .models import Member
 
 logger = logging.getLogger(__name__)
 
+WaffleFlag = get_waffle_flag_model()
+
 
 class HttpResponseTemporaryRedirect(HttpResponseRedirect):
     """
@@ -93,6 +99,28 @@ def __call__(self, request):
         return get_production_redirect(request)
 
 
+class CustomWaffleMiddleware(WaffleMiddleware):
+    """
+    Add testing flags to request.waffle_tests.
+
+    Parent middleware process_response saves waffles and waffle_tests to a
+    cookie in the response.
+    """
+
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        for flag in WaffleFlag.objects.filter(testing=True):
+            tc = get_waffle_setting("TEST_COOKIE") % flag.name
+            if tc in request.GET:
+                on = request.GET[tc] == "1"
+                if not hasattr(request, "waffle_tests"):
+                    request.waffle_tests = {}
+                request.waffle_tests[flag.name] = on
+        return self.process_response(request, self.get_response(request))
+
+
 class AddMemberMiddleware(object):
     """
     A convenience middleware that adds the Member to the request if the user is

open_humans/settings.py

@@ -207,7 +207,7 @@ def to_bool(env, default="false"):
     "django.contrib.messages.middleware.MessageMiddleware",
     "django.middleware.clickjacking.XFrameOptionsMiddleware",
     "open_humans.middleware.AddMemberMiddleware",
-    "waffle.middleware.WaffleMiddleware",
+    "open_humans.middleware.CustomWaffleMiddleware",
     "django.middleware.cache.FetchFromCacheMiddleware",
 )