From 63ad073efd0b20f59f02bc37182c0180cac3f405 Mon Sep 17 00:00:00 2001
From: Michael Gundlach <michael.gundlach@rackspace.com>
Date: Wed, 15 Sep 2010 09:25:53 -0400
Subject: RateLimitingMiddleware tests

---
 nova/api/rackspace/__init__.py           | 24 ++++++++++++++++--------
 nova/api/rackspace/ratelimiting/tests.py |  3 +++
 2 files changed, 19 insertions(+), 8 deletions(-)

(limited to 'nova/api')

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