import datetime
import json
import logging
import os
import shutil
import moto
import py
import pytest
from .. import lambda_handler
class TestCreateStacks:
@pytest.fixture
def setup(self):
with moto.mock_cloudformation(), moto.mock_dynamodb2(), moto.mock_sts():
setup = lambda_handler.Setup(logging)
yield setup
def test_create_stacks_deployment(self, setup):
# make mock_rules directory
os.mkdir("auto_remediate_setup/data/mock_rules")
# move mock CloudFormation Stacks
shutil.copyfile(
"auto_remediate_setup/test/data/mock_rules/cloudtrail-enabled.json",
"auto_remediate_setup/data/mock_rules/cloudtrail-enabled.json",
)
# backup settings
shutil.move(
"auto_remediate_setup/data/auto-remediate-settings.json",
"auto_remediate_setup/data/auto-remediate-settings-backup.json",
)
# move mock settings
shutil.copyfile(
"auto_remediate_setup/test/data/auto-remediate-settings-deploy.json",
"auto_remediate_setup/data/auto-remediate-settings.json",
)
# create table
setup.client_dynamodb.create_table(
TableName="settings-table",
KeySchema=[{"AttributeName": "key", "KeyType": "HASH"}],
AttributeDefinitions=[{"AttributeName": "key", "AttributeType": "S"}],
ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
)
# insert settings to DynamoDB
os.environ["SETTINGSTABLE"] = "settings-table"
setup.setup_dynamodb()
# test create_stacks function
setup.create_stacks("mock_rules", setup.get_settings())
# validate Stack created
response = setup.client_cloudformation.list_stacks()
assert response["StackSummaries"][0]["StackName"] == "cloudtrail-enabled"
# restore settings
shutil.move(
"auto_remediate_setup/data/auto-remediate-settings-backup.json",
"auto_remediate_setup/data/auto-remediate-settings.json",
)
# delete mock_rules directory
shutil.rmtree("auto_remediate_setup/data/mock_rules")
# def test_create_stacks_removal(self, setup):
# # make mock_rules directory
# os.mkdir("auto_remediate_setup/data/mock_rules")
# # move mock CloudFormation Stacks
# shutil.copyfile(
# "auto_remediate_setup/test/data/mock_rules/cloudtrail-enabled.json",
# "auto_remediate_setup/data/mock_rules/cloudtrail-enabled.json",
# )
# # backup settings
# shutil.move(
# "auto_remediate_setup/data/auto-remediate-settings.json",
# "auto_remediate_setup/data/auto-remediate-settings-backup.json",
# )
# # move mock deploy settings
# shutil.copyfile(
# "auto_remediate_setup/test/data/auto-remediate-settings-deploy.json",
# "auto_remediate_setup/data/auto-remediate-settings.json",
# )
# # create table
# setup.client_dynamodb.create_table(
# TableName="settings-table",
# KeySchema=[{"AttributeName": "key", "KeyType": "HASH"}],
# AttributeDefinitions=[{"AttributeName": "key", "AttributeType": "S"}],
# ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
# )
# # insert settings to DynamoDB
# os.environ["SETTINGSTABLE"] = "settings-table"
# setup.setup_dynamodb()
# # test create_stacks function
# setup.create_stacks("mock_rules", setup.get_settings())
# # validate Stack created
# response = setup.client_cloudformation.list_stacks()
# assert response["StackSummaries"][0]["StackName"] == "cloudtrail-enabled"
# # move mock remove settings
# shutil.copyfile(
# "auto_remediate_setup/test/data/auto-remediate-settings-remove.json",
# "auto_remediate_setup/data/auto-remediate-settings.json",
# )
# # insert settings to DynamoDB
# setup.setup_dynamodb()
# # test create_stacks function
# setup.create_stacks("mock_rules", setup.get_settings())
# # validate Stack created
# response = setup.client_cloudformation.list_stacks()
# # restore settings
# shutil.move(
# "auto_remediate_setup/data/auto-remediate-settings-backup.json",
# "auto_remediate_setup/data/auto-remediate-settings.json",
# )
# # delete mock_rules directory
# shutil.rmtree("auto_remediate_setup/data/mock_rules")
class TestGetSettings:
@pytest.fixture
def setup(self):
with moto.mock_cloudformation(), moto.mock_dynamodb2(), moto.mock_sts():
setup = lambda_handler.Setup(logging)
yield setup
def test_get_settings(self, setup):
"""Tests retrieval of settings from DynamoDB
Arguments:
setup {class} -- Instance of Setup class
"""
os.environ["SETTINGSTABLE"] = "settings-table"
# create table
setup.client_dynamodb.create_table(
TableName="settings-table",
KeySchema=[{"AttributeName": "key", "KeyType": "HASH"}],
AttributeDefinitions=[{"AttributeName": "key", "AttributeType": "S"}],
ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
)
# populate table
setup.client_dynamodb.put_item(
TableName="settings-table",
Item={"key": {"S": "version"}, "value": {"N": "1.0"}},
)
# test get_settings function
settings = setup.get_settings()
# validate test
assert settings["version"] == 1.0
def test_invalid_table_schema(self, setup):
"""Tests retrieval of settings from DynamoDB with the wrong schema
Arguments:
setup {class} -- Instance of Setup class
"""
os.environ["SETTINGSTABLE"] = "settings-table"
setup.client_dynamodb.create_table(
TableName="settings-table",
KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}],
AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}],
ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
)
setup.client_dynamodb.put_item(
TableName="settings-table", Item={"id": {"S": "123"}}
)
# test get_settings function
assert setup.get_settings() == {}
def test_no_table_name(self, setup):
"""Tests retrieval of settings from DynamoDB with no table name
Arguments:
setup {class} -- Instance of Setup class
"""
assert setup.get_settings() == {}
class TestGetCurrentStacks:
@pytest.fixture
def setup(self):
with moto.mock_cloudformation(), moto.mock_dynamodb2(), moto.mock_sts():
setup = lambda_handler.Setup(logging)
yield setup
def test_get_current_stacks(self, setup):
"""Tests retrieval of CloudFormation Stacks
Arguments:
setup {[type]} -- [description]
"""
setup.client_cloudformation.create_stack(
StackName="sample_sqs",
TemplateBody='{"Resources":{"SQSQueue":{"Type":"AWS::SQS::Queue","Properties":{"QueueName":"test_queue"}}}}',
)
# test get_current_stacks function
response = setup.get_current_stacks()
# assert stack created
assert response[0] == "sample_sqs"
def test_no_stacks(self, setup):
"""Tests retrieval of CloudFormation Stacks with no Stacks
Arguments:
setup {[type]} -- [description]
"""
assert setup.get_current_stacks() == []
class TestSetupDynamoDb:
@pytest.fixture
def setup(self):
with moto.mock_dynamodb2(), moto.mock_sts():
setup = lambda_handler.Setup(logging)
yield setup
def test_setup_dynamodb(self, setup):
os.environ["SETTINGSTABLE"] = "settings-table"
# create table
setup.client_dynamodb.create_table(
TableName="settings-table",
KeySchema=[{"AttributeName": "key", "KeyType": "HASH"}],
AttributeDefinitions=[{"AttributeName": "key", "AttributeType": "S"}],
ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
)
setup.setup_dynamodb()
assert len(setup.client_dynamodb.scan(TableName="settings-table")["Items"]) > 0
