From ffb521f86b2dc793e0c4a5bc953e1ae7aadb5195 Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Sat, 16 May 2020 15:03:26 +0100 Subject: [PATCH] Linting --- moto/athena/models.py | 7 ++- moto/athena/responses.py | 45 +++++++-------- tests/test_athena/test_athena.py | 94 ++++++++++++++++++-------------- 3 files changed, 78 insertions(+), 68 deletions(-) diff --git a/moto/athena/models.py b/moto/athena/models.py index 20d180d74..c39c13817 100644 --- a/moto/athena/models.py +++ b/moto/athena/models.py @@ -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 diff --git a/moto/athena/responses.py b/moto/athena/responses.py index c572cea0b..b52e0beed 100644 --- a/moto/athena/responses.py +++ b/moto/athena/responses.py @@ -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), ) diff --git a/tests/test_athena/test_athena.py b/tests/test_athena/test_athena.py index 597361b1d..93ca436aa 100644 --- a/tests/test_athena/test_athena.py +++ b/tests/test_athena/test_athena.py @@ -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'] - 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/' - } - }, - 'Description': 'Test work group' - }) + 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/"} + }, + "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/",} + }, )