Add tests

.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.

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/

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)
 

requirements.txt

@@ -7,3 +7,4 @@ awslimitchecker
 nose
 Flask-Script
 sendgrid
+pytest

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
+