Exclude busy users from roster suggest

Also, ignore W503/W504 in flake8 checks to allow line breaks
before or after binary operators

e2e/test_roster_suggest.py

@@ -80,4 +80,53 @@ def test_api_v0_fill_gap(user, team, role, roster, event):
     re = requests.get(api_v0('teams/%s/rosters/%s/%s/suggest?start=%s' %
                              (team_name, roster_name, role_name, start + 2000)))
     assert re.status_code == 200
+    assert re.json()['user'] == user_name
+
+@prefix('test_v0_fill_gap_skip_busy')
+def test_api_v0_fill_gap_skip_busy(user, team, role, roster, event):
+    user_name = user.create()
+    user_name_2 = user.create()
+    user_name_3 = user.create()
+    user_name_4 = user.create()
+    team_name = team.create()
+    role_name = role.create()
+    role_name_2 = role.create()
+    roster_name = roster.create(team_name)
+    start = int(time.time()) + 1000
+    user.add_to_roster(user_name, team_name, roster_name)
+    user.add_to_roster(user_name_2, team_name, roster_name)
+    user.add_to_roster(user_name_3, team_name, roster_name)
+    user.add_to_roster(user_name_4, team_name, roster_name)
+
+    # Create events: user_name will be the expected user, with events far from
+    # the suggestion time (start + 2000). user_name_4 will be a busy user, who
+    # would otherwise be chosen.
+    event.create({'start': start,
+                  'end': start + 1000,
+                  'user': user_name,
+                  'team': team_name,
+                  'role': role_name})
+    event.create({'start': start + 1000,
+                  'end': start + 2000,
+                  'user': user_name_2,
+                  'team': team_name,
+                  'role': role_name})
+    event.create({'start': start + 3000,
+                  'end': start + 4000,
+                  'user': user_name_3,
+                  'team': team_name,
+                  'role': role_name})
+    event.create({'start': start + 4000,
+                  'end': start + 5000,
+                  'user': user_name,
+                  'team': team_name,
+                  'role': role_name})
+    event.create({'start': start + 1000,
+                  'end': start + 3000,
+                  'user': user_name_4,
+                  'team': team_name,
+                  'role': role_name_2})
+    re = requests.get(api_v0('teams/%s/rosters/%s/%s/suggest?start=%s' %
+                             (team_name, roster_name, role_name, start + 2000)))
+    assert re.status_code == 200
     assert re.json()['user'] == user_name

src/oncall/api/v0/roster_suggest.py

@@ -14,22 +14,30 @@ def on_get(req, resp, team, roster, role):
             raise HTTPBadRequest('Invalid role')
         role_id = cursor.fetchone()[0]
 
-        cursor.execute('SELECT id FROM team WHERE name = %s', team)
-        if cursor.rowcount == 0:
-            raise HTTPBadRequest('Invalid team')
-        team_id = cursor.fetchone()[0]
-
-        cursor.execute('SELECT id FROM roster WHERE name = %s and team_id = %s', (roster, team_id))
+        cursor.execute('''SELECT `team`.`id`, `roster`.`id` FROM `team` JOIN `roster` ON `roster`.`team_id` = `team`.`id`
+                          WHERE `roster`.`name` = %s and `team`.`name` = %s''', (roster, team))
         if cursor.rowcount == 0:
             raise HTTPBadRequest('Invalid roster')
-        roster_id = cursor.fetchone()[0]
+        team_id, roster_id = cursor.fetchone()
 
         cursor.execute('SELECT COUNT(*) FROM roster_user WHERE roster_id = %s', roster_id)
         if cursor.rowcount == 0:
             raise HTTPNotFound()
         roster_size = cursor.fetchone()[0]
         length = 604800 * roster_size
 
+        data = {'team_id': team_id,
+                'roster_id': roster_id,
+                'role_id': role_id,
+                'past': start - length,
+                'start': start,
+                'future': start + length}
+
+        cursor.execute('''SELECT `user`.`name` FROM `event` JOIN `user` ON `event`.`user_id` = `user`.`id`
+                          WHERE `team_id` = %(team_id)s AND %(start)s BETWEEN `event`.`start` AND `event`.`end`''',
+                       data)
+        busy_users = set(row[0] for row in cursor)
+
         cursor.execute('''SELECT * FROM
                             (SELECT `user`.`name` AS `user`, MAX(`event`.`start`) AS `before`
                              FROM `roster_user` JOIN `user` ON `user`.`id` = `roster_user`.`user_id`
@@ -45,25 +53,24 @@ def on_get(req, resp, team, roster, role):
                                AND `role_id` = %(role_id)s AND `start` BETWEEN %(start)s AND %(future)s
                              GROUP BY `user`.`name`) future
                           USING (`user`)''',
-                       {'team_id': team_id,
-                        'roster_id': roster_id,
-                        'role_id': role_id,
-                        'past': start - length,
-                        'start': start,
-                        'future': start + length})
+                       data)
         candidate = None
         max_score = -1
         # Find argmax(min(time between start and last event, time before start and next event))
         # If no next/last event exists, set value to infinity
         # This should maximize gaps between shifts
+        ret = {}
         for (user, before, after) in cursor:
+            if user in busy_users:
+                continue
             before = start - before if before is not None else float('inf')
             after = after - start if after is not None else float('inf')
             score = min(before, after)
+            ret[user] = score if score != float('inf') else 'infinity'
             if score > max_score:
                 candidate = user
                 max_score = score
     finally:
         cursor.close()
         connection.close()
-    resp.body = json_dumps({'user': candidate})
+    resp.body = json_dumps({'user': candidate, 'data': ret})

tox.ini

@@ -4,7 +4,7 @@ envlist = py27
 [flake8]
 exclude = .tox,./build
 filename = *.py
-ignore = E101,E741,W292,E722,E731
+ignore = E101,E741,W292,E722,E731,W503,W504
 max-line-length = 160
 
 [testenv]