diff --git a/.gitignore b/.gitignore index 908a064..09e464e 100644 --- a/.gitignore +++ b/.gitignore @@ -23,7 +23,7 @@ var/ *.egg-info/ .installed.cfg *.egg - +.idea/* # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. diff --git a/Makefile b/Makefile index 94dabee..6d408dc 100644 --- a/Makefile +++ b/Makefile @@ -8,8 +8,7 @@ run: @python awslimits/server.py test: - rm -f .coverage - nosetests $(NOSE_ARGS) ./tests/ + @py.test tests/ travis: nosetests --with-coverage ./tests/ diff --git a/awslimits/support.py b/awslimits/support.py index d12f79d..c5c61e7 100644 --- a/awslimits/support.py +++ b/awslimits/support.py @@ -6,9 +6,9 @@ import time from datetime import datetime, timedelta from decimal import Decimal +from botocore.exceptions import ClientError import settings -from dynamo_helpers import create_or_get_table TICKETS_TABLE_NAME = 'awslimits_tickets' LIMITS_TABLE_NAME = 'awslimits_limits' @@ -21,6 +21,7 @@ def dict_to_obj(dict_): struct = namedtuple('struct', dict_.keys()) return struct(**dict_) + def get_boto_resource(resource): sts = boto3.client('sts') assumed_role = sts.assume_role(RoleArn=settings.ROLE_ARN, RoleSessionName="awslimits") @@ -52,6 +53,30 @@ def get_boto_client(client, region_name=settings.REGION_NAME): ) +def create_or_get_table(table_name, attribute_definitions, key_schema): + dynamodb = get_boto_resource('dynamodb') + + try: + table = dynamodb.create_table( + AttributeDefinitions=attribute_definitions, + TableName=table_name, + KeySchema=key_schema, + ProvisionedThroughput={ + 'ReadCapacityUnits': 1, + 'WriteCapacityUnits': 1, + }, + ) + except ClientError as exc: + if exc.response['Error']['Code'] == 'ResourceInUseException': + table = dynamodb.Table(table_name) + return table + else: + raise + + table.meta.client.get_waiter('table_exists').wait(TableName=table_name) + return table + + def get_aws_limit_checker(): return AwsLimitChecker(region=settings.REGION_NAME, account_id=settings.ACCOUNT_ID, account_role=settings.ACCOUNT_ROLE) diff --git a/requirements.txt b/requirements.txt index d4b58f5..fbb9534 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ awslimitchecker nose Flask-Script sendgrid +pytest diff --git a/tests/unit/test_limits.py b/tests/unit/test_limits.py new file mode 100644 index 0000000..db628f3 --- /dev/null +++ b/tests/unit/test_limits.py @@ -0,0 +1,88 @@ +import moto +import boto3 +from unittest import TestCase +from mock import patch +import os + +mock_env = { + 'ROLE_ARN': 'arn:aws:iam::212311232123:role/test', + 'SENDGRID_API_KEY': '1231231231231231231', + 'AWS_ACCESS_KEY_ID': 'awskey', + 'AWS_SECRET_ACCESS_KEY': 'awssecret', +} +with patch.dict(os.environ, mock_env): + from awslimits.support import create_or_get_table + + +@moto.mock_dynamodb2 +@moto.mock_sts +class TestDynamo(TestCase): + def test_create_or_get_new_table(self): + dynamodb = boto3.client('dynamodb', region_name="us-east-1") + table_name = 'awslimits_limits' + initial_table_list = dynamodb.list_tables() + assert table_name not in initial_table_list["TableNames"] + table = create_or_get_table( + table_name='awslimits_limits', + attribute_definitions=[ + { + 'AttributeName': 'limit_name', + 'AttributeType': 'S', + }, + ], + key_schema=[ + { + 'AttributeName': 'limit_name', + 'KeyType': 'HASH' + }, + ], + ) + assert table.name == 'awslimits_limits' + final_table_list = dynamodb.list_tables() + assert table_name in final_table_list["TableNames"] + + + def test_create_or_get_exiting_table(self): + dynamodb = boto3.client('dynamodb', region_name="us-east-1") + table_name = 'awslimits_limits1' + initial_table_list = dynamodb.list_tables() + assert table_name not in initial_table_list["TableNames"] + table = create_or_get_table( + table_name='awslimits_limits1', + attribute_definitions=[ + { + 'AttributeName': 'limit_name', + 'AttributeType': 'S', + }, + ], + key_schema=[ + { + 'AttributeName': 'limit_name', + 'KeyType': 'HASH' + }, + ], + ) + table = create_or_get_table( + table_name='awslimits_limits1', + attribute_definitions=[ + { + 'AttributeName': 'limit_name', + 'AttributeType': 'S', + }, + ], + key_schema=[ + { + 'AttributeName': 'limit_name', + 'KeyType': 'HASH' + }, + ], + ) + final_table_list = dynamodb.list_tables() + assert table_name in final_table_list["TableNames"] + + def test_create_or_get_table_exception(self): + pass + + def test_create_or_get_table_failed(self): + pass +