Merge "use get_container_info in ratelimit"

2 files changed, 26 insertions, 27 deletions

diff --git a/swift/common/middleware/ratelimit.py b/swift/common/middleware/ratelimit.py
index 6e8dd8729..06823cd77 100644
--- a/swift/common/middleware/ratelimit.py
+++ b/swift/common/middleware/ratelimit.py
@@ -18,8 +18,7 @@ from swift import gettext_ as _
 import eventlet
 
 from swift.common.utils import cache_from_env, get_logger, register_swift_info
-from swift.proxy.controllers.base import get_container_memcache_key, \
-    get_account_info
+from swift.proxy.controllers.base import get_account_info, get_container_info
 from swift.common.memcached import MemcacheConnectionError
 from swift.common.swob import Request, Response
 
@@ -118,11 +117,10 @@ class RateLimitMiddleware(object):
         self.container_listing_ratelimits = interpret_conf_limits(
             conf, 'container_listing_ratelimit_')
 
-    def get_container_size(self, account_name, container_name):
+    def get_container_size(self, env):
         rv = 0
-        memcache_key = get_container_memcache_key(account_name,
-                                                  container_name)
-        container_info = self.memcache_client.get(memcache_key)
+        container_info = get_container_info(
+            env, self.app, swift_source='RL')
         if isinstance(container_info, dict):
             rv = container_info.get(
                 'object_count', container_info.get('container_size', 0))
@@ -149,8 +147,7 @@ class RateLimitMiddleware(object):
 
         if account_name and container_name and obj_name and \
                 req.method in ('PUT', 'DELETE', 'POST', 'COPY'):
-            container_size = self.get_container_size(
-                account_name, container_name)
+            container_size = self.get_container_size(req.environ)
             container_rate = get_maxrate(
                 self.container_ratelimits, container_size)
             if container_rate:
@@ -160,8 +157,7 @@ class RateLimitMiddleware(object):
 
         if account_name and container_name and not obj_name and \
                 req.method == 'GET':
-            container_size = self.get_container_size(
-                account_name, container_name)
+            container_size = self.get_container_size(req.environ)
             container_rate = get_maxrate(
                 self.container_listing_ratelimits, container_size)
             if container_rate:
diff --git a/test/unit/common/middleware/test_ratelimit.py b/test/unit/common/middleware/test_ratelimit.py
index a502bafc9..12940eac8 100644
--- a/test/unit/common/middleware/test_ratelimit.py
+++ b/test/unit/common/middleware/test_ratelimit.py
@@ -189,7 +189,7 @@ class TestRateLimit(unittest.TestCase):
         the_app = ratelimit.filter_factory(conf_dict)(FakeApp())
         the_app.memcache_client = fake_memcache
         req = lambda: None
-        req.environ = {}
+        req.environ = {'swift.cache': fake_memcache, 'PATH_INFO': '/v1/a/c/o'}
         with mock.patch('swift.common.middleware.ratelimit.get_account_info',
                         lambda *args, **kwargs: {}):
             req.method = 'DELETE'
@@ -243,7 +243,7 @@ class TestRateLimit(unittest.TestCase):
         the_app.memcache_client = fake_memcache
         req = lambda: None
         req.method = 'PUT'
-        req.environ = {}
+        req.environ = {'PATH_INFO': '/v1/a/c/o', 'swift.cache': fake_memcache}
         with mock.patch('swift.common.middleware.ratelimit.get_account_info',
                         lambda *args, **kwargs: {}):
             tuples = the_app.get_ratelimitable_key_tuples(req, 'a', 'c', 'o')
@@ -255,20 +255,23 @@ class TestRateLimit(unittest.TestCase):
         conf_dict = {'account_ratelimit': current_rate}
         self.test_ratelimit = ratelimit.filter_factory(conf_dict)(FakeApp())
         ratelimit.http_connect = mock_http_connect(204)
-        with mock.patch('swift.common.middleware.ratelimit.get_account_info',
+        with mock.patch('swift.common.middleware.ratelimit.get_container_info',
                         lambda *args, **kwargs: {}):
-            for meth, exp_time in [
-                    ('DELETE', 9.8), ('GET', 0), ('POST', 0), ('PUT', 9.8)]:
-                req = Request.blank('/v/a%s/c' % meth)
-                req.method = meth
-                req.environ['swift.cache'] = FakeMemcache()
-                make_app_call = lambda: self.test_ratelimit(req.environ,
-                                                            start_response)
-                begin = time.time()
-                self._run(make_app_call, num_calls, current_rate,
-                          check_time=bool(exp_time))
-                self.assertEquals(round(time.time() - begin, 1), exp_time)
-                self._reset_time()
+            with mock.patch(
+                    'swift.common.middleware.ratelimit.get_account_info',
+                    lambda *args, **kwargs: {}):
+                for meth, exp_time in [('DELETE', 9.8), ('GET', 0),
+                                       ('POST', 0), ('PUT', 9.8)]:
+                    req = Request.blank('/v/a%s/c' % meth)
+                    req.method = meth
+                    req.environ['swift.cache'] = FakeMemcache()
+                    make_app_call = lambda: self.test_ratelimit(req.environ,
+                                                                start_response)
+                    begin = time.time()
+                    self._run(make_app_call, num_calls, current_rate,
+                              check_time=bool(exp_time))
+                    self.assertEquals(round(time.time() - begin, 1), exp_time)
+                    self._reset_time()
 
     def test_ratelimit_set_incr(self):
         current_rate = 5
@@ -403,7 +406,7 @@ class TestRateLimit(unittest.TestCase):
         req.method = 'PUT'
         req.environ['swift.cache'] = FakeMemcache()
         req.environ['swift.cache'].set(
-            ratelimit.get_container_memcache_key('a', 'c'),
+            get_container_memcache_key('a', 'c'),
             {'container_size': 1})
 
         time_override = [0, 0, 0, 0, None]
@@ -437,7 +440,7 @@ class TestRateLimit(unittest.TestCase):
         req.method = 'GET'
         req.environ['swift.cache'] = FakeMemcache()
         req.environ['swift.cache'].set(
-            ratelimit.get_container_memcache_key('a', 'c'),
+            get_container_memcache_key('a', 'c'),
             {'container_size': 1})
 
         time_override = [0, 0, 0, 0, None]