APIGateway: Added API for usage plans

moto/apigateway/models.py

@@ -309,6 +309,19 @@ class ApiKey(BaseModel, dict):
         self['stageKeys'] = stageKeys
 
 
+class UsagePlan(BaseModel, dict):
+
+    def __init__(self, name=None, description=None, apiStages=[],
+                 throttle=None, quota=None):
+        super(UsagePlan, self).__init__()
+        self['id'] = create_id()
+        self['name'] = name
+        self['description'] = description
+        self['apiStages'] = apiStages
+        self['throttle'] = throttle
+        self['quota'] = quota
+
+
 class RestAPI(BaseModel):
 
     def __init__(self, id, region_name, name, description):
@@ -408,6 +421,7 @@ class APIGatewayBackend(BaseBackend):
         super(APIGatewayBackend, self).__init__()
         self.apis = {}
         self.keys = {}
+        self.usage_plans = {}
         self.region_name = region_name
 
     def reset(self):
@@ -576,6 +590,21 @@ class APIGatewayBackend(BaseBackend):
         self.keys.pop(api_key_id)
         return {}
 
+    def create_usage_plan(self, payload):
+        plan = UsagePlan(**payload)
+        self.usage_plans[plan['id']] = plan
+        return plan
+
+    def get_usage_plans(self):
+        return list(self.usage_plans.values())
+
+    def get_usage_plan(self, usage_plan_id):
+        return self.usage_plans[usage_plan_id]
+
+    def delete_usage_plan(self, usage_plan_id):
+        self.usage_plans.pop(usage_plan_id)
+        return {}
+
 
 apigateway_backends = {}
 for region_name in Session().get_available_regions('apigateway'):

moto/apigateway/responses.py

@@ -248,3 +248,25 @@ class APIGatewayResponse(BaseResponse):
         elif self.method == 'DELETE':
             apikey_response = self.backend.delete_apikey(apikey)
         return 200, {}, json.dumps(apikey_response)
+
+    def usage_plans(self, request, full_url, headers):
+        self.setup_class(request, full_url, headers)
+
+        if self.method == 'POST':
+            usage_plan_response = self.backend.create_usage_plan(json.loads(self.body))
+        elif self.method == 'GET':
+            usage_plans_response = self.backend.get_usage_plans()
+            return 200, {}, json.dumps({"item": usage_plans_response})
+        return 200, {}, json.dumps(usage_plan_response)
+
+    def usage_plan_individual(self, request, full_url, headers):
+        self.setup_class(request, full_url, headers)
+
+        url_path_parts = self.path.split("/")
+        usage_plan = url_path_parts[2]
+
+        if self.method == 'GET':
+            usage_plan_response = self.backend.get_usage_plan(usage_plan)
+        elif self.method == 'DELETE':
+            usage_plan_response = self.backend.delete_usage_plan(usage_plan)
+        return 200, {}, json.dumps(usage_plan_response)

moto/apigateway/urls.py

@@ -20,4 +20,6 @@ url_paths = {
     '{0}/restapis/(?P<function_id>[^/]+)/resources/(?P<resource_id>[^/]+)/methods/(?P<method_name>[^/]+)/integration/responses/(?P<status_code>\d+)/?$': APIGatewayResponse().integration_responses,
     '{0}/apikeys$': APIGatewayResponse().apikeys,
     '{0}/apikeys/(?P<apikey>[^/]+)': APIGatewayResponse().apikey_individual,
+    '{0}/usageplans$': APIGatewayResponse().usage_plans,
+    '{0}/usageplans/(?P<usage_plan_id>[^/]+)': APIGatewayResponse().usage_plan_individual,
 }

tests/test_apigateway/test_apigateway.py

@@ -995,3 +995,40 @@ def test_api_keys():
 
     response = client.get_api_keys()
     len(response['items']).should.equal(1)
+
+@mock_apigateway
+def test_usage_plans():
+    region_name = 'us-west-2'
+    client = boto3.client('apigateway', region_name=region_name)
+    response = client.get_usage_plans()
+    len(response['items']).should.equal(0)
+
+    usage_plan_name = 'TEST-PLAN'
+    payload = {'name': usage_plan_name}
+    response = client.create_usage_plan(**payload)
+    usage_plan = client.get_usage_plan(usagePlanId=response['id'])
+    usage_plan['name'].should.equal(usage_plan_name)
+    usage_plan['apiStages'].should.equal([])
+
+    usage_plan_name = 'TEST-PLAN-2'
+    usage_plan_description = 'Description'
+    usage_plan_quota = {'limit': 10, 'period': 'DAY', 'offset': 0}
+    usage_plan_throttle = {'rateLimit': 2, 'burstLimit': 1}
+    usage_plan_api_stages = [{'apiId': 'foo', 'stage': 'bar'}]
+    payload = {'name': usage_plan_name, 'description': usage_plan_description, 'quota': usage_plan_quota, 'throttle': usage_plan_throttle, 'apiStages': usage_plan_api_stages}
+    response = client.create_usage_plan(**payload)
+    usage_plan_id = response['id']
+    usage_plan = client.get_usage_plan(usagePlanId=usage_plan_id)
+    usage_plan['name'].should.equal(usage_plan_name)
+    usage_plan['description'].should.equal(usage_plan_description)
+    usage_plan['apiStages'].should.equal(usage_plan_api_stages)
+    usage_plan['throttle'].should.equal(usage_plan_throttle)
+    usage_plan['quota'].should.equal(usage_plan_quota)
+
+    response = client.get_usage_plans()
+    len(response['items']).should.equal(2)
+
+    client.delete_usage_plan(usagePlanId=usage_plan_id)
+
+    response = client.get_usage_plans()
+    len(response['items']).should.equal(1)

tests/test_apigateway/test_server.py

@@ -1,5 +1,6 @@
 from __future__ import unicode_literals
 import sure  # noqa
+import json
 
 import moto.server as server
 
@@ -9,8 +10,50 @@ Test the different server responses
 
 
 def test_list_apis():
-    backend = server.create_backend_app("apigateway")
+    backend = server.create_backend_app('apigateway')
     test_client = backend.test_client()
 
     res = test_client.get('/restapis')
     res.data.should.equal(b'{"item": []}')
+
+def test_usage_plans_apis():
+    backend = server.create_backend_app('apigateway')
+    test_client = backend.test_client()
+
+    '''
+    List usage plans (expect empty)
+    '''
+    res = test_client.get('/usageplans')
+    json.loads(res.data)["item"].should.have.length_of(0)
+
+    '''
+    Create usage plan
+    '''
+    res = test_client.post('/usageplans', data=json.dumps({'name': 'test'}))
+    created_plan = json.loads(res.data)
+    created_plan['name'].should.equal('test')
+
+    '''
+    List usage plans (expect 1 plan)
+    '''
+    res = test_client.get('/usageplans')
+    json.loads(res.data)["item"].should.have.length_of(1)
+
+    '''
+    Get single usage plan
+    '''
+    res = test_client.get('/usageplans/{0}'.format(created_plan["id"]))
+    fetched_plan = json.loads(res.data)
+    fetched_plan.should.equal(created_plan)
+
+    '''
+    Delete usage plan
+    '''
+    res = test_client.delete('/usageplans/{0}'.format(created_plan["id"]))
+    res.data.should.equal(b'{}')
+
+    '''
+    List usage plans (expect empty again)
+    '''
+    res = test_client.get('/usageplans')
+    json.loads(res.data)["item"].should.have.length_of(0)