Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Linting

Browse Source
This commit is contained in:
Bert Blommers 2020-05-16 15:03:26 +01:00
parent dd20fec9f3
commit ffb521f86b
3 changed files with 78 additions and 68 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
moto/athena/models.py
View File

@ -50,7 +50,6 @@ class WorkGroup(TaggableResourceMixin, BaseModel):
class Execution(BaseModel):
def __init__(self, query, context, config, workgroup):
self.id = str(uuid4())
self.query = query
@ -97,11 +96,13 @@ class AthenaBackend(BaseBackend):
"State": wg.state,
"Configuration": wg.configuration,
"Description": wg.description,
"CreationTime": time.time()
"CreationTime": time.time(),
}
def start_query_execution(self, query, context, config, workgroup):
execution = Execution(query=query, context=context, config=config, workgroup=workgroup)
execution = Execution(
query=query, context=context, config=config, workgroup=workgroup
)
self.executions[execution.id] = execution
return execution.id

45
moto/athena/responses.py
View File

@ -43,32 +43,34 @@ class AthenaResponse(BaseResponse):
workgroup = self._get_param("WorkGroup")
if workgroup and not self.athena_backend.get_work_group(workgroup):
return self.error("WorkGroup does not exist", 400)
id = self.athena_backend.start_query_execution(query=query, context=context, config=config, workgroup=workgroup)
id = self.athena_backend.start_query_execution(
query=query, context=context, config=config, workgroup=workgroup
)
return json.dumps({"QueryExecutionId": id})
def get_query_execution(self):
exec_id = self._get_param("QueryExecutionId")
execution = self.athena_backend.get_execution(exec_id)
result = {
'QueryExecution': {
'QueryExecutionId': exec_id,
'Query': execution.query,
'StatementType': 'DDL',
'ResultConfiguration': execution.config,
'QueryExecutionContext': execution.context,
'Status': {
'State': execution.status,
'SubmissionDateTime': execution.start_time
"QueryExecution": {
"QueryExecutionId": exec_id,
"Query": execution.query,
"StatementType": "DDL",
"ResultConfiguration": execution.config,
"QueryExecutionContext": execution.context,
"Status": {
"State": execution.status,
"SubmissionDateTime": execution.start_time,
},
'Statistics': {
'EngineExecutionTimeInMillis': 0,
'DataScannedInBytes': 0,
'TotalExecutionTimeInMillis': 0,
'QueryQueueTimeInMillis': 0,
'QueryPlanningTimeInMillis': 0,
'ServiceProcessingTimeInMillis': 0
"Statistics": {
"EngineExecutionTimeInMillis": 0,
"DataScannedInBytes": 0,
"TotalExecutionTimeInMillis": 0,
"QueryQueueTimeInMillis": 0,
"QueryPlanningTimeInMillis": 0,
"ServiceProcessingTimeInMillis": 0,
},
'WorkGroup': execution.workgroup
"WorkGroup": execution.workgroup,
}
}
return json.dumps(result)
@ -80,11 +82,6 @@ class AthenaResponse(BaseResponse):
def error(self, msg, status):
return (
json.dumps(
{
"__type": "InvalidRequestException",
"Message": msg,
}
),
json.dumps({"__type": "InvalidRequestException", "Message": msg,}),
dict(status=status),
)

90
tests/test_athena/test_athena.py
View File

@ -64,18 +64,18 @@ def test_create_and_get_workgroup():
create_basic_workgroup(client=client, name="athena_workgroup")
work_group = client.get_work_group(WorkGroup='athena_workgroup')['WorkGroup']
work_group = client.get_work_group(WorkGroup="athena_workgroup")["WorkGroup"]
del work_group["CreationTime"] # Were not testing creationtime atm
work_group.should.equal({
'Name': 'athena_workgroup',
'State': 'ENABLED',
'Configuration': {
'ResultConfiguration': {
'OutputLocation': 's3://bucket-name/prefix/'
}
work_group.should.equal(
{
"Name": "athena_workgroup",
"State": "ENABLED",
"Configuration": {
"ResultConfiguration": {"OutputLocation": "s3://bucket-name/prefix/"}
},
'Description': 'Test work group'
})
"Description": "Test work group",
}
)
@mock_athena
@ -83,16 +83,20 @@ def test_start_query_execution():
client = boto3.client("athena", region_name="us-east-1")
create_basic_workgroup(client=client, name="athena_workgroup")
response = client.start_query_execution(QueryString='query1',
QueryExecutionContext={'Database': 'string'},
ResultConfiguration={'OutputLocation': 'string'},
WorkGroup='athena_workgroup')
assert 'QueryExecutionId' in response
response = client.start_query_execution(
QueryString="query1",
QueryExecutionContext={"Database": "string"},
ResultConfiguration={"OutputLocation": "string"},
WorkGroup="athena_workgroup",
)
assert "QueryExecutionId" in response
sec_response = client.start_query_execution(QueryString='query2',
QueryExecutionContext={'Database': 'string'},
ResultConfiguration={'OutputLocation': 'string'})
assert 'QueryExecutionId' in sec_response
sec_response = client.start_query_execution(
QueryString="query2",
QueryExecutionContext={"Database": "string"},
ResultConfiguration={"OutputLocation": "string"},
)
assert "QueryExecutionId" in sec_response
response["QueryExecutionId"].shouldnt.equal(sec_response["QueryExecutionId"])
@ -101,10 +105,12 @@ def test_start_query_validate_workgroup():
client = boto3.client("athena", region_name="us-east-1")
with assert_raises(ClientError) as err:
client.start_query_execution(QueryString='query1',
QueryExecutionContext={'Database': 'string'},
ResultConfiguration={'OutputLocation': 'string'},
WorkGroup='unknown_workgroup')
client.start_query_execution(
QueryString="query1",
QueryExecutionContext={"Database": "string"},
ResultConfiguration={"OutputLocation": "string"},
WorkGroup="unknown_workgroup",
)
err.exception.response["Error"]["Code"].should.equal("InvalidRequestException")
err.exception.response["Error"]["Message"].should.equal("WorkGroup does not exist")
@ -117,9 +123,11 @@ def test_get_query_execution():
location = "s3://bucket-name/prefix/"
database = "database"
# Start Query
exex_id = client.start_query_execution(QueryString=query,
QueryExecutionContext={'Database': database},
ResultConfiguration={'OutputLocation': location})["QueryExecutionId"]
exex_id = client.start_query_execution(
QueryString=query,
QueryExecutionContext={"Database": database},
ResultConfiguration={"OutputLocation": location},
)["QueryExecutionId"]
#
details = client.get_query_execution(QueryExecutionId=exex_id)["QueryExecution"]
#
@ -129,12 +137,16 @@ def test_get_query_execution():
details["ResultConfiguration"]["OutputLocation"].should.equal(location)
details["QueryExecutionContext"]["Database"].should.equal(database)
details["Status"]["State"].should.equal("QUEUED")
details["Statistics"].should.equal({'EngineExecutionTimeInMillis': 0,
'DataScannedInBytes': 0,
'TotalExecutionTimeInMillis': 0,
'QueryQueueTimeInMillis': 0,
'QueryPlanningTimeInMillis': 0,
'ServiceProcessingTimeInMillis': 0})
details["Statistics"].should.equal(
{
"EngineExecutionTimeInMillis": 0,
"DataScannedInBytes": 0,
"TotalExecutionTimeInMillis": 0,
"QueryQueueTimeInMillis": 0,
"QueryPlanningTimeInMillis": 0,
"ServiceProcessingTimeInMillis": 0,
}
)
assert "WorkGroup" not in details
@ -146,9 +158,11 @@ def test_stop_query_execution():
location = "s3://bucket-name/prefix/"
database = "database"
# Start Query
exex_id = client.start_query_execution(QueryString=query,
QueryExecutionContext={'Database': database},
ResultConfiguration={'OutputLocation': location})["QueryExecutionId"]
exex_id = client.start_query_execution(
QueryString=query,
QueryExecutionContext={"Database": database},
ResultConfiguration={"OutputLocation": location},
)["QueryExecutionId"]
# Stop Query
client.stop_query_execution(QueryExecutionId=exex_id)
# Verify status
@ -163,8 +177,6 @@ def create_basic_workgroup(client, name):
Name=name,
Description="Test work group",
Configuration={
"ResultConfiguration": {
"OutputLocation": "s3://bucket-name/prefix/",
}
}
"ResultConfiguration": {"OutputLocation": "s3://bucket-name/prefix/",}
},
)