RateLimitingMiddleware tests

3 files changed, 52 insertions, 26 deletions

diff --git a/nova/api/rackspace/__init__.py b/nova/api/rackspace/__init__.py
index e35109b43..66d80a5b7 100644
--- a/nova/api/rackspace/__init__.py
+++ b/nova/api/rackspace/__init__.py
@@ -92,23 +92,31 @@ class RateLimitingMiddleware(wsgi.Middleware):
         action_name = self.get_action_name(req)
         if not action_name: # not rate limited
             return self.application
-        delay = self.limiter.perform(action_name, username=username)
-        if action_name == 'POST servers':
-            # "POST servers" is a POST, so it counts against "POST" too.
-            delay2 = self.limiter.perform('POST', username=username)
-            delay = max(delay or 0, delay2 or 0)
+        delay = self.get_delay(action_name, username)
         if delay:
             # TODO(gundlach): Get the retry-after format correct.
             raise webob.exc.HTTPRequestEntityTooLarge(headers={
                     'Retry-After': time.time() + delay})
-        else:
-            return self.application
+        return self.application
+
+    def get_delay(self, action_name, username):
+        """Return the delay for the given action and username, or None if
+        the action would not be rate limited.
+        """
+        if action_name == 'POST servers':
+            # "POST servers" is a POST, so it counts against "POST" too.
+            # Attempt the "POST" first, lest we are rate limited by "POST" but
+            # use up a precious "POST servers" call.
+            delay = self.limiter.perform("POST", username=username)
+            if delay:
+                return delay
+        return self.limiter.perform(action_name, username=username)
 
     def get_action_name(self, req):
         """Return the action name for this request."""
         if req.method == 'GET' and 'changes-since' in req.GET:
             return 'GET changes-since'
-        if req.method == 'POST' and req.path_info.starts_with('/servers'):
+        if req.method == 'POST' and req.path_info.startswith('/servers'):
             return 'POST servers'
         if req.method in ['PUT', 'POST', 'DELETE']:
             return req.method
diff --git a/nova/api/rackspace/ratelimiting/tests.py b/nova/api/rackspace/ratelimiting/tests.py
index 1983cdea8..545e1d1b6 100644
--- a/nova/api/rackspace/ratelimiting/tests.py
+++ b/nova/api/rackspace/ratelimiting/tests.py
@@ -58,3 +58,6 @@ class Test(unittest.TestCase):
         self.exhaust('c', 0, username='chuck')
         self.exhaust('c', 0, username='bob')
         self.exhaust('c', 0, username='alice')
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/nova/tests/api/rackspace/__init__.py b/nova/tests/api/rackspace/__init__.py
index f7537a4e7..2fab1a4da 100644
--- a/nova/tests/api/rackspace/__init__.py
+++ b/nova/tests/api/rackspace/__init__.py
@@ -17,22 +17,15 @@
 
 import unittest
 
-from nova.api.rackspace.ratelimiting import RateLimitingMiddleware
+from nova.api.rackspace import RateLimitingMiddleware
 from nova.tests.api.test_helper import *
 from webob import Request
 
 class RateLimitingMiddlewareTest(unittest.TestCase):
-    def setUp(self):
-        self.middleware = RateLimitingMiddleware(APIStub())
-        self.stubs = stubout.StubOutForTesting()
-
-    def tearDown(self):
-        self.stubs.UnsetAll()
-
     def test_get_action_name(self):
         middleware = RateLimitingMiddleware(APIStub())
         def verify(method, url, action_name):
-            req = Request(url)
+            req = Request.blank(url)
             req.method = method
             action = middleware.get_action_name(req)
             self.assertEqual(action, action_name)
@@ -45,12 +38,34 @@ class RateLimitingMiddlewareTest(unittest.TestCase):
         verify('GET', '/servers/4', None)
         verify('HEAD', '/servers/4', None)
 
-    def TODO_test_call(self):
-        pass
-        #mw = make_middleware()
-        #req = build_request('DELETE', '/servers/4')
-        #for i in range(5):
-        #    resp = req.get_response(mw)
-        #    assert resp is OK
-        #resp = req.get_response(mw)
-        #assert resp is rate limited
+    def exhaust(self, middleware, method, url, username, times):
+        req = Request.blank(url, dict(REQUEST_METHOD=method),
+                            headers={'X-Auth-User': username})
+        for i in range(times):
+            resp = req.get_response(middleware)
+            self.assertEqual(resp.status_int, 200)
+        resp = req.get_response(middleware)
+        self.assertEqual(resp.status_int, 413)
+        self.assertTrue('Retry-After' in resp.headers)
+
+    def test_single_action(self):
+        middleware = RateLimitingMiddleware(APIStub())
+        self.exhaust(middleware, 'DELETE', '/servers/4', 'usr1', 100)
+        self.exhaust(middleware, 'DELETE', '/servers/4', 'usr2', 100)
+
+    def test_POST_servers_action_implies_POST_action(self):
+        middleware = RateLimitingMiddleware(APIStub())
+        self.exhaust(middleware, 'POST', '/servers/4', 'usr1', 10)
+        self.exhaust(middleware, 'POST', '/images/4', 'usr2', 10)
+        self.assertTrue(set(middleware.limiter._levels) == 
+                        set(['usr1:POST', 'usr1:POST servers', 'usr2:POST']))
+
+    def test_POST_servers_action_correctly_ratelimited(self):
+        middleware = RateLimitingMiddleware(APIStub())
+        # Use up all of our "POST" allowance for the minute, 5 times
+        for i in range(5):
+            self.exhaust(middleware, 'POST', '/servers/4', 'usr1', 10)
+            # Reset the 'POST' action counter.
+            del middleware.limiter._levels['usr1:POST']
+        # All 50 daily "POST servers" actions should be all used up
+        self.exhaust(middleware, 'POST', '/servers/4', 'usr1', 0)