Issue 2276 - Handle groups in notification.json requests

sitescripts/notifications/web/notification.py

 # coding: utf-8
 
 # This file is part of the Adblock Plus web scripts,
 # Copyright (C) 2006-2015 Eyeo GmbH
 #
 # Adblock Plus is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 3 as
 # published by the Free Software Foundation.
 #
 # Adblock Plus is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU General Public License for more details.
 #
 # You should have received a copy of the GNU General Public License
 # along with Adblock Plus.  If not, see <http://www.gnu.org/licenses/>.
 
+import copy
 import json
 import random
 import time
-from urlparse import parse_qs
+import urlparse
 
 from sitescripts.notifications.parser import load_notifications
 from sitescripts.web import url_handler
 
+def _determine_groups(version, notifications):
+  version_groups = dict(x.split("/") for x in version.split("-")[1:]
+                        if x.count("/") == 1)
+  groups = []
+  for notification in notifications:
+    if "variants" not in notification:
+      continue
+    group_id = notification["id"]
+    if group_id in version_groups:
+      groups.append({"id": group_id, "variant": int(version_groups[group_id])})
+  return groups
+
 def _assign_groups(notifications):
   groups = []
   selection = random.random()
-  base = 0
-  for notification in filter(lambda notification: "variants" in notification,
-                             notifications):
+  start = 0
+  for notification in notifications:
+    if "variants" not in notification:
+      continue
     group = {"id": notification["id"], "variant": 0}
     groups.append(group)
     for i, variant in enumerate(notification["variants"]):
       sample_size = variant["sample"]
-      if (group["variant"] == 0 and sample_size > 0 and selection >= base and
-          selection <= sample_size + base):
+      end = start + sample_size
+      selected = sample_size > 0 and start <= selection <= end
+      start = end
+      if selected:
         group["variant"] = i + 1
-      base += sample_size
+        break
   return groups
 
-def _create_response(notifications, groups):
-  response = {
-    "version": time.strftime("%Y%m%d%H%M", time.gmtime())
-  }
+def _get_active_variant(notifications, groups):
   for group in groups:
     group_id = group["id"]
     variant = group["variant"]
-    response["version"] += "-%s/%s" % (group_id, variant)
-    if variant > 0:
-      notification = next(notification for notification in notifications
-                          if notification["id"] == group_id)
-      notification = notification.copy()
-      notification.update(notification["variants"][variant - 1])
-      del notification["variants"]
-      response["notifications"] = [notification]
-  if not "notifications" in response:
-    response["notifications"] = \
-      filter(lambda notification: not "variants" in notification, notifications)
+    if variant == 0:
+      continue
+    notification = next(x for x in notifications if x["id"] == group_id)
+    notification = copy.deepcopy(notification)
+    notification.update(notification["variants"][variant - 1])
+    for key_to_remove in ("sample", "variants"):
+      notification.pop(key_to_remove, None)
+    return notification
+
+def _generate_version(groups):
+  version = time.strftime("%Y%m%d%H%M", time.gmtime())
+  for group in groups:
+    version += "-%s/%s" % (group["id"], group["variant"])
+  return version
+
+def _create_response(notifications, groups):
+  active_variant = _get_active_variant(notifications, groups)
+  if active_variant:
+    notifications = [active_variant]
+  else:
+    notifications = [x for x in notifications if "variants" not in x]
+  response = {
+    "version": _generate_version(groups),
+    "notifications": notifications
+  }
   return json.dumps(response, ensure_ascii=False, indent=2,
                     separators=(",", ": "), sort_keys=True)
 
 @url_handler("/notification.json")
 def notification(environ, start_response):
-  params = parse_qs(environ.get("QUERY_STRING", ""))
+  params = urlparse.parse_qs(environ.get("QUERY_STRING", ""))
   version = params.get("lastVersion", [""])[0]
-  current_groups = dict(component.split("/")
-                        for component in version.split("-")[1:]
-                        if "/" in component)
   notifications = load_notifications()
-  groups = []
-  for notification in notifications:
-    if "variants" in notification:
-      group_id = notification["id"]
-      if group_id in current_groups:
-        groups.append({"id": group_id,
-                       "variant": int(current_groups[group_id])})
+  groups = _determine_groups(version, notifications)
   if not groups:
     groups = _assign_groups(notifications)
   response = _create_response(notifications, groups)
-  start_response("200 OK", ["Content-Type", "application/json"])
-  return response
+  start_response("200 OK",
+                 [("Content-Type", "application/json; charset=utf-8")])
+  return response.encode("utf-8")