Merge pull request #2996 from bblommers/athena_start_stop_execution

Athena - Start/Stop query execution, get_work_group

tests/test_athena/test_athena.py

@@ -1,8 +1,7 @@
 from __future__ import unicode_literals
 
-import datetime
-
 from botocore.exceptions import ClientError
+from nose.tools import assert_raises
 import boto3
 import sure  # noqa
 
@@ -57,3 +56,127 @@ def test_create_work_group():
     work_group["Name"].should.equal("athena_workgroup")
     work_group["Description"].should.equal("Test work group")
     work_group["State"].should.equal("ENABLED")
+
+
+@mock_athena
+def test_create_and_get_workgroup():
+    client = boto3.client("athena", region_name="us-east-1")
+
+    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",
+        }
+    )
+
+
+@mock_athena
+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
+
+    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"])
+
+
+@mock_athena
+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",
+        )
+    err.exception.response["Error"]["Code"].should.equal("InvalidRequestException")
+    err.exception.response["Error"]["Message"].should.equal("WorkGroup does not exist")
+
+
+@mock_athena
+def test_get_query_execution():
+    client = boto3.client("athena", region_name="us-east-1")
+
+    query = "SELECT stuff"
+    location = "s3://bucket-name/prefix/"
+    database = "database"
+    # Start Query
+    exex_id = client.start_query_execution(
+        QueryString=query,
+        QueryExecutionContext={"Database": database},
+        ResultConfiguration={"OutputLocation": location},
+    )["QueryExecutionId"]
+    #
+    details = client.get_query_execution(QueryExecutionId=exex_id)["QueryExecution"]
+    #
+    details["QueryExecutionId"].should.equal(exex_id)
+    details["Query"].should.equal(query)
+    details["StatementType"].should.equal("DDL")
+    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,
+        }
+    )
+    assert "WorkGroup" not in details
+
+
+@mock_athena
+def test_stop_query_execution():
+    client = boto3.client("athena", region_name="us-east-1")
+
+    query = "SELECT stuff"
+    location = "s3://bucket-name/prefix/"
+    database = "database"
+    # Start Query
+    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
+    details = client.get_query_execution(QueryExecutionId=exex_id)["QueryExecution"]
+    #
+    details["QueryExecutionId"].should.equal(exex_id)
+    details["Status"]["State"].should.equal("CANCELLED")
+
+
+def create_basic_workgroup(client, name):
+    client.create_work_group(
+        Name=name,
+        Description="Test work group",
+        Configuration={
+            "ResultConfiguration": {"OutputLocation": "s3://bucket-name/prefix/",}
+        },
+    )