93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
import os
|
|
from functools import wraps
|
|
from uuid import uuid4
|
|
|
|
import boto3
|
|
|
|
from moto import mock_aws
|
|
|
|
|
|
def lakeformation_aws_verified(func):
|
|
"""
|
|
Function that is verified to work against AWS.
|
|
Can be run against AWS at any time by setting:
|
|
MOTO_TEST_ALLOW_AWS_REQUEST=true
|
|
|
|
If this environment variable is not set, the function runs in a `mock_aws` context.
|
|
|
|
Note that LakeFormation is not enabled by default - visit the AWS Console to permit access to the user who executes these tests.
|
|
"""
|
|
|
|
@wraps(func)
|
|
def pagination_wrapper():
|
|
bucket_name = str(uuid4())
|
|
|
|
allow_aws_request = (
|
|
os.environ.get("MOTO_TEST_ALLOW_AWS_REQUEST", "false").lower() == "true"
|
|
)
|
|
|
|
if allow_aws_request:
|
|
resp = create_glue_infra_and_test(bucket_name)
|
|
else:
|
|
with mock_aws():
|
|
resp = create_glue_infra_and_test(bucket_name)
|
|
return resp
|
|
|
|
def create_glue_infra_and_test(bucket_name):
|
|
glue = boto3.client("glue", region_name="eu-west-2")
|
|
lf = boto3.client("lakeformation", region_name="eu-west-2")
|
|
s3 = boto3.client("s3", region_name="us-east-1")
|
|
|
|
s3.create_bucket(Bucket=bucket_name)
|
|
s3.put_bucket_tagging(
|
|
Bucket=bucket_name,
|
|
Tagging={"TagSet": [{"Key": "environment", "Value": "moto_tests"}]},
|
|
)
|
|
lf.register_resource(
|
|
ResourceArn=f"arn:aws:s3:::{bucket_name}", UseServiceLinkedRole=True
|
|
)
|
|
|
|
db_name = str(uuid4())[0:6]
|
|
table_name = str(uuid4())[0:6]
|
|
column_name = str(uuid4())[0:6]
|
|
glue.create_database(
|
|
DatabaseInput={"Name": db_name}, Tags={"environment": "moto_tests"}
|
|
)
|
|
glue.create_table(
|
|
DatabaseName=db_name,
|
|
TableInput={
|
|
"Name": table_name,
|
|
"StorageDescriptor": {
|
|
"Columns": [{"Name": column_name, "Type": "string"}]
|
|
},
|
|
},
|
|
)
|
|
|
|
try:
|
|
resp = func(bucket_name, db_name, table_name, column_name)
|
|
finally:
|
|
### CLEANUP ###
|
|
|
|
glue.delete_table(DatabaseName=db_name, Name=table_name)
|
|
glue.delete_database(Name=db_name)
|
|
|
|
lf.deregister_resource(ResourceArn=f"arn:aws:s3:::{bucket_name}")
|
|
|
|
versions = s3.list_object_versions(Bucket=bucket_name).get("Versions", [])
|
|
for key in versions:
|
|
s3.delete_object(
|
|
Bucket=bucket_name, Key=key["Key"], VersionId=key.get("VersionId")
|
|
)
|
|
delete_markers = s3.list_object_versions(Bucket=bucket_name).get(
|
|
"DeleteMarkers", []
|
|
)
|
|
for key in delete_markers:
|
|
s3.delete_object(
|
|
Bucket=bucket_name, Key=key["Key"], VersionId=key.get("VersionId")
|
|
)
|
|
s3.delete_bucket(Bucket=bucket_name)
|
|
|
|
return resp
|
|
|
|
return pagination_wrapper
|