| #!/usr/bin/env python |
| # |
| # Copyright 2010 Google Inc. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| """QuotaService tests.""" |
| |
| __author__ = 'rafek@google.com (Rafe Kaplan)' |
| |
| import unittest |
| import time |
| |
| from protorpc import test_util |
| |
| from quotas import services |
| |
| SIMPLE_CONFIG = services.QuotaConfig( |
| buckets=[ |
| services.QuotaConfig.Bucket( |
| name='DISK', |
| initial_tokens=1000), |
| services.QuotaConfig.Bucket( |
| name='EMAILS', |
| initial_tokens=10, |
| refresh_every=10) |
| ]) |
| |
| |
| class QuotaTestBase(test_util.TestCase): |
| |
| def MakeState(self, disk=None, emails=None, next_refresh=2000, state=None): |
| state = state or {} |
| if disk: |
| state['DISK'] = disk, None |
| if emails: |
| state['EMAILS'] = emails, next_refresh |
| buckets = dict((bucket.name, bucket) for bucket in SIMPLE_CONFIG.buckets) |
| return services.QuotaState(state, buckets) |
| |
| def setUp(self): |
| self.current_time = 1000.5 |
| self.original_time = time.time |
| time.time = self.Time |
| |
| def tearDown(self): |
| time.time = self.original_time |
| |
| def Time(self): |
| return self.current_time |
| |
| |
| class QuotaStateTest(QuotaTestBase): |
| |
| def setUp(self): |
| super(QuotaStateTest, self).setUp() |
| self.state = self.MakeState() |
| |
| def testGetRemainingTokens(self): |
| self.state.begin_transaction() |
| self.assertEquals(1000, self.state.get_remaining_tokens('DISK')) |
| self.assertEquals(10, self.state.get_remaining_tokens('EMAILS')) |
| self.assertEquals(None, self.state.get_remaining_tokens('UNKNOWN')) |
| self.state.abort_transaction() |
| |
| self.state = self.MakeState(200, 5) |
| self.state.begin_transaction() |
| self.assertEquals(200, self.state.get_remaining_tokens('DISK')) |
| self.assertEquals(5, self.state.get_remaining_tokens('EMAILS')) |
| self.assertEquals(None, self.state.get_remaining_tokens('UNKNOWN')) |
| self.state.abort_transaction() |
| |
| def testCheckQuota(self): |
| self.state.begin_transaction() |
| self.assertEquals(100, self.state.check_quota('DISK', 100)) |
| self.assertEquals(1000, self.state.check_quota('DISK', 1000)) |
| self.assertEquals(1000, self.state.check_quota('DISK', 1001)) |
| self.assertEquals(-1001, self.state.check_quota('DISK', -1001)) |
| self.assertEquals(1, self.state.check_quota('EMAILS', 1)) |
| self.assertEquals(10, self.state.check_quota('EMAILS', 10)) |
| self.assertEquals(10, self.state.check_quota('EMAILS', 11)) |
| self.assertEquals(-11, self.state.check_quota('EMAILS', -11)) |
| self.assertEquals(None, self.state.check_quota('UNKNOWN', 10)) |
| self.state.abort_transaction() |
| |
| def testDeduct(self): |
| self.state.begin_transaction() |
| self.assertEquals(100, self.state.deduct_quota('DISK', 100)) |
| self.assertEquals(900, self.state.get_remaining_tokens('DISK')) |
| self.state.commit_transaction() |
| |
| self.state.begin_transaction() |
| self.assertEquals(900, self.state.deduct_quota('DISK', 900)) |
| self.assertEquals(0, self.state.get_remaining_tokens('DISK')) |
| self.state.commit_transaction() |
| |
| self.state.begin_transaction() |
| self.assertEquals(0, self.state.deduct_quota('DISK', 1)) |
| self.assertEquals(0, self.state.get_remaining_tokens('DISK')) |
| self.state.commit_transaction() |
| |
| self.state.begin_transaction() |
| self.assertEquals(10, self.state.deduct_quota('EMAILS', 11)) |
| self.assertEquals(0, self.state.get_remaining_tokens('EMAILS')) |
| self.state.commit_transaction() |
| |
| self.state.begin_transaction() |
| self.assertEquals(0, self.state.deduct_quota('DISK', 1)) |
| self.assertEquals(0, self.state.get_remaining_tokens('EMAILS')) |
| self.state.commit_transaction() |
| |
| def testDeductAbort(self): |
| self.state.begin_transaction() |
| self.assertEquals(100, self.state.deduct_quota('DISK', 100)) |
| self.assertEquals(900, self.state.get_remaining_tokens('DISK')) |
| self.state.abort_transaction() |
| |
| self.state.begin_transaction() |
| self.assertEquals(1000, self.state.get_remaining_tokens('DISK')) |
| self.state.abort_transaction() |
| |
| def testAutoRefresh(self): |
| # Use some quota. |
| self.state.begin_transaction() |
| self.state.deduct_quota('DISK', 100) |
| self.state.deduct_quota('EMAILS', 10) |
| self.state.commit_transaction() |
| |
| # Cause refresh. |
| self.current_time += 20 |
| |
| # Emails should now be restored. |
| self.state.begin_transaction() |
| self.assertEquals(900, self.state.get_remaining_tokens('DISK')) |
| self.assertEquals(10, self.state.get_remaining_tokens('EMAILS')) |
| self.state.commit_transaction() |
| |
| def testAutoRefreshDuringTransaction(self): |
| # Use some quota quota. |
| self.state.begin_transaction() |
| self.state.deduct_quota('DISK', 100) |
| self.state.deduct_quota('EMAILS', 10) |
| |
| # Refresh is too late for transaction. |
| self.current_time += 20 |
| |
| self.assertEquals(900, self.state.get_remaining_tokens('DISK')) |
| self.assertEquals(0, self.state.get_remaining_tokens('EMAILS')) |
| self.state.commit_transaction() |
| |
| def testDeductSameBucket(self): |
| self.state.begin_transaction() |
| self.state.deduct_quota('DISK', 100) |
| self.assertRaises(AssertionError, |
| self.state.deduct_quota, 'DISK', 100) |
| self.state.commit_transaction() |
| |
| |
| class QuotaServiceTest(QuotaTestBase): |
| |
| def setUp(self): |
| super(QuotaServiceTest, self).setUp() |
| self.states = {} |
| self.service = services.QuotaService(SIMPLE_CONFIG, {}) |
| |
| def MakeCheckQuotaRequest(self, user, quotas): |
| return services.QuotaRequest( |
| user=user, |
| quotas=[ |
| services.QuotaCheck(name=bucket, |
| tokens=tokens, |
| mode=services.QuotaCheck.Mode(mode)) |
| for bucket, tokens, mode in quotas |
| ]) |
| |
| def AssertRemaining(self, user, remaining_quotas): |
| response = self.service.get_quota_state( |
| services.QuotaStateRequest(user=user)) |
| self.assertEquals(len(remaining_quotas), len(response.bucket_states)) |
| for (name, remaining), bucket_state in zip(remaining_quotas, |
| response.bucket_states): |
| self.assertEquals(name, bucket_state.name) |
| self.assertEquals(remaining, bucket_state.remaining_tokens) |
| |
| def testCheckQuotaAll(self): |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 100, 'ALL'), |
| ('EMAILS', 10, 'ALL'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.OK, response.all_status) |
| self.assertFalse(response.denied) |
| self.assertEquals(2, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.OK, response.results[0].status) |
| self.assertEquals(100, response.results[0].available) |
| |
| self.assertEquals( |
| services.CheckResult.Status.OK, response.results[1].status) |
| self.assertEquals(10, response.results[1].available) |
| |
| self.AssertRemaining('bob', [('DISK', 900), ('EMAILS', 0)]) |
| |
| def testCheckQuotaAll_SomeAvailable(self): |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 100, 'ALL'), |
| ('EMAILS', 11, 'ALL'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.SOME, response.all_status) |
| self.assertTrue(response.denied) |
| self.assertEquals(2, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.OK, response.results[0].status) |
| self.assertEquals(100, response.results[0].available) |
| |
| self.assertEquals( |
| services.CheckResult.Status.SOME, response.results[1].status) |
| self.assertEquals(10, response.results[1].available) |
| |
| self.AssertRemaining('bob', [('DISK', 1000), ('EMAILS', 10)]) |
| |
| def testCheckQuotaAll_NoneAvailable(self): |
| # Deplete all quota. |
| self.service.check_quota(self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1000, 'ALL'), |
| ('EMAILS', 10, 'ALL'), |
| ])) |
| |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1, 'ALL'), |
| ('EMAILS', 1, 'ALL'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.NONE, response.all_status) |
| self.assertTrue(response.denied) |
| self.assertEquals(2, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.NONE, response.results[0].status) |
| self.assertEquals(0, response.results[0].available) |
| |
| self.assertEquals( |
| services.CheckResult.Status.NONE, response.results[1].status) |
| self.assertEquals(0, response.results[1].available) |
| |
| self.AssertRemaining('bob', [('DISK', 0), ('EMAILS', 0)]) |
| |
| def testCheckQuotaAll_SomeAndNone(self): |
| # Deplete some quota. |
| self.service.check_quota(self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1000, 'ALL'), |
| ])) |
| |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1, 'ALL'), |
| ('EMAILS', 11, 'ALL'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.SOME, response.all_status) |
| self.assertTrue(response.denied) |
| self.assertEquals(2, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.NONE, response.results[0].status) |
| self.assertEquals(0, response.results[0].available) |
| |
| self.assertEquals( |
| services.CheckResult.Status.SOME, response.results[1].status) |
| self.assertEquals(10, response.results[1].available) |
| |
| self.AssertRemaining('bob', [('DISK', 0), ('EMAILS', 10)]) |
| |
| def testCheckQuotaSome(self): |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 500, 'ALL'), |
| ('EMAILS', 5, 'SOME'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.OK, response.all_status) |
| self.assertFalse(response.denied) |
| self.assertEquals(2, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.OK, response.results[0].status) |
| self.assertEquals(500, response.results[0].available) |
| |
| self.assertEquals( |
| services.CheckResult.Status.OK, response.results[1].status) |
| self.assertEquals(5, response.results[1].available) |
| |
| self.AssertRemaining('bob', [('DISK', 500), ('EMAILS', 5)]) |
| |
| def testCheckQuotaSome_SomeAvailable(self): |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 100, 'SOME'), |
| ('EMAILS', 11, 'SOME'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.SOME, response.all_status) |
| self.assertFalse(response.denied) |
| self.assertEquals(2, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.OK, response.results[0].status) |
| self.assertEquals(100, response.results[0].available) |
| |
| self.assertEquals( |
| services.CheckResult.Status.SOME, response.results[1].status) |
| self.assertEquals(10, response.results[1].available) |
| |
| self.AssertRemaining('bob', [('DISK', 900), ('EMAILS', 0)]) |
| |
| def testCheckQuotaSome_NoneAvailable(self): |
| # Deplete all quota. |
| self.service.check_quota(self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1000, 'ALL'), |
| ('EMAILS', 10, 'ALL'), |
| ])) |
| |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1, 'SOME'), |
| ('EMAILS', 1, 'SOME'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.NONE, response.all_status) |
| self.assertTrue(response.denied) |
| self.assertEquals(2, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.NONE, response.results[0].status) |
| self.assertEquals(0, response.results[0].available) |
| |
| self.assertEquals( |
| services.CheckResult.Status.NONE, response.results[1].status) |
| self.assertEquals(0, response.results[1].available) |
| |
| self.AssertRemaining('bob', [('DISK', 0), ('EMAILS', 0)]) |
| |
| def testCheckQuotaSome_SomeAndNone(self): |
| # Deplete some quota. |
| self.service.check_quota(self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1000, 'ALL'), |
| ])) |
| |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1, 'SOME'), |
| ('EMAILS', 11, 'SOME'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.SOME, response.all_status) |
| self.assertTrue(response.denied) |
| self.assertEquals(2, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.NONE, response.results[0].status) |
| self.assertEquals(0, response.results[0].available) |
| |
| self.assertEquals( |
| services.CheckResult.Status.SOME, response.results[1].status) |
| self.assertEquals(10, response.results[1].available) |
| |
| self.AssertRemaining('bob', [('DISK', 0), ('EMAILS', 10)]) |
| |
| def testCheckQuotaCheckAll(self): |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 500, 'CHECK_ALL'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.OK, response.all_status) |
| self.assertFalse(response.denied) |
| self.assertEquals(1, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.OK, response.results[0].status) |
| self.assertEquals(500, response.results[0].available) |
| |
| self.AssertRemaining('bob', [('DISK', 1000), ('EMAILS', 10)]) |
| |
| def testCheckQuotaCheckAll_SomeAvailable(self): |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1001, 'CHECK_ALL'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.SOME, response.all_status) |
| self.assertTrue(response.denied) |
| self.assertEquals(1, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.SOME, response.results[0].status) |
| self.assertEquals(1000, response.results[0].available) |
| |
| self.AssertRemaining('bob', [('DISK', 1000), ('EMAILS', 10)]) |
| |
| def testCheckQuotaCheckAll_NoneAvailable(self): |
| # Deplete all quota. |
| self.service.check_quota(self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1000, 'ALL'), |
| ])) |
| |
| request = self.MakeCheckQuotaRequest('bob', |
| [('DISK', 1, 'CHECK_ALL'), |
| ]) |
| response = self.service.check_quota(request) |
| |
| self.assertEquals(services.CheckResult.Status.NONE, response.all_status) |
| self.assertTrue(response.denied) |
| self.assertEquals(1, len(response.results)) |
| |
| self.assertEquals( |
| services.CheckResult.Status.NONE, response.results[0].status) |
| self.assertEquals(0, response.results[0].available) |
| |
| self.AssertRemaining('bob', [('DISK', 0), ('EMAILS', 10)]) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
