build rate limiting: tests

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):

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()
 

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()
-