From eb69abff8b0b24d3269d49aa15b11010bf3a8442 Mon Sep 17 00:00:00 2001
From: Jimmy Zelinskie <jimmy.zelinskie+git@gmail.com>
Date: Tue, 6 Dec 2016 14:47:02 -0500
Subject: [PATCH] build rate limiting: tests

---
 data/queue.py         | 10 +++++++---
 endpoints/building.py |  9 ++++++---
 test/test_queue.py    | 19 ++++++++++++++++++-
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/data/queue.py b/data/queue.py
index 5cd84726d..659a485b0 100644
--- a/data/queue.py
+++ b/data/queue.py
@@ -77,10 +77,14 @@ class WorkQueue(object):
             ._available_jobs(now, name_match_query)
             .where(~(QueueItem.queue_name << running_query)))
 
-  def num_available_jobs(self, available_min_time, prefix):
+  def num_available_jobs_between(self, available_min_time, available_max_time, prefix):
+    """
+    Returns the number of available queue items with a given prefix between the two provided times.
+    """
     prefix = prefix.lstrip('/')
-    available = self._available_jobs(datetime.utcnow(),
-                                     self._name_match_query() + prefix)
+    available = self._available_jobs(available_max_time,
+                                     '/'.join([self._queue_name, prefix]) + '%')
+
     return available.where(QueueItem.available_after >= available_min_time).count()
 
   def _name_match_query(self):
diff --git a/endpoints/building.py b/endpoints/building.py
index a94154efa..d2684c8b8 100644
--- a/endpoints/building.py
+++ b/endpoints/building.py
@@ -25,10 +25,13 @@ class MaximumBuildsQueuedException(Exception):
 
 
 def start_build(repository, prepared_build, pull_robot_name=None):
-  queue_item_prefix = '%s/%s' % repository.namespace_user.username, repository.name
+  queue_item_prefix = '%s/%s' % (repository.namespace_user.username, repository.name)
   if MAX_BUILD_QUEUE_ITEMS > 0 and MAX_BUILD_QUEUE_SECS > 0:
-    available_min = datetime.utcnow() - timedelta(seconds=-MAX_BUILD_QUEUE_SECS)
-    available_builds = dockerfile_build_queue.num_available_jobs(available_min, queue_item_prefix)
+    now = datetime.utcnow()
+    available_min = now - timedelta(seconds=MAX_BUILD_QUEUE_SECS)
+    available_builds = dockerfile_build_queue.num_available_jobs_between(available_min,
+                                                                         now,
+                                                                         queue_item_prefix)
     if available_builds > MAX_BUILD_QUEUE_ITEMS:
       raise MaximumBuildsQueuedException()
 
diff --git a/test/test_queue.py b/test/test_queue.py
index d31268850..0ac6cb633 100644
--- a/test/test_queue.py
+++ b/test/test_queue.py
@@ -2,6 +2,7 @@ import unittest
 import json
 import time
 
+from datetime import datetime, timedelta
 from functools import wraps
 
 from app import app
@@ -182,6 +183,22 @@ class TestQueue(QueueTestCase):
       self.assertIn(msg, seen)
 
 
+  def test_num_available_between(self):
+    now = datetime.utcnow()
+    self.queue.put(['abc', 'def'], self.TEST_MESSAGE_1, available_after=-10)
+    self.queue.put(['abc', 'ghi'], self.TEST_MESSAGE_2, available_after=-5)
+
+    # Partial results
+    count = self.queue.num_available_jobs_between(now-timedelta(seconds=8), now, 'abc')
+    self.assertEqual(1, count)
+
+    # All results
+    count = self.queue.num_available_jobs_between(now-timedelta(seconds=20), now, 'abc')
+    self.assertEqual(2, count)
+
+    # No results
+    count = self.queue.num_available_jobs_between(now, now, 'abc')
+    self.assertEqual(0, count)
+
 if __name__ == '__main__':
   unittest.main()
-