使用 Moto 进行 Pytest，使用后端更改 athena 查询的状态

Question

我正在使用 moto 来测试我的代码库中的 aws 功能。我遇到的问题之一是，在测试 athena 时，查询状态无限期地保持在“QUEUED”状态，导致测试失败或超时。

下面是要测试的方法：

import time
import boto3

class Athena:

    CLIENT = boto3.client("athena")

    class QueryError(Exception):
        """A class for exceptions related to queries."""

    @classmethod
    def execute_query(cls, query, result_location, check_status=True, 
    time_limit=10):
        """
        Execute a query in Athena.
        """

        _result_configuration = {"OutputLocation": result_location}

        _kwargs = {"QueryString": query, "ResultConfiguration": 
        _result_configuration}

        response = cls.CLIENT.start_query_execution(**_kwargs)
        query_id = response["QueryExecutionId"]

        if check_status:
            old_time = time.time()
            while True:
                status = cls.CLIENT.get_query_execution(
                QueryExecutionId=query_id)
                status = status["QueryExecution"]["Status"]["State"]
                    if status in ["SUCCEEDED", "FAILED", "CANCELLED"]:
                        if status == "FAILED":
                            raise cls.QueryError("error")
                        break
                    time.sleep(0.2)  # 200ms

                    if time.time() - old_time > time_limit and status 
                    == "QUEUED":
                        raise cls.QueryError("time limit reached")

        return query_id

这是通过测试的夹具

from moto.s3 import mock_s3
import boto3

@pytest.fixture
def s3():
    with mock_s3():
        s3 = boto3.client("s3")
        yield s3

这是测试（记住你需要from x用上面的方法换成模块）

import uuid
import boto3
import pytest
from moto.athena import mock_athena
from moto.s3 import mock_s3


@mock_s3
@mock_athena
def test_execute_query_check(s3):
    from x import Athena

    """
    Test for 'execute_query' (with status check)
    """
    CLIENT = s3
    bucket_name = "pytest." + str(uuid.uuid4())

    # Bucket creation
    bucket_config = {"LocationConstraint": "us-east-2"}
    CLIENT.create_bucket(Bucket=bucket_name, 
    CreateBucketConfiguration=bucket_config)
    waiter = CLIENT.get_waiter("bucket_exists")
    waiter.wait(Bucket=bucket_name)

    s3_location = f"s3://{bucket_name}/"
    query = "SELECT current_date, current_time;"
    query_id = Athena.execute_query(query, s3_location, 
    check_status=True)
    assert query_id

此测试失败，因为moto不会更改过去查询的状态"QUEUED"，并且测试期望更改为状态，否则会触发异常。

我希望能够做类似的事情：

from moto.athena import athena_backends
athena_backends['us-east-2'].job_flows[query_id].state = "SUCCEEDED"

正如本期所建议的：https ://github.com/spulec/moto/issues/380

然而，“作业流”属性在 boto3 mapreduce 后端似乎不再存在，而且我找不到显式更改它的方法。理想情况下，这可以在测试中的某个地方发生，以手动更改查询的状态以模拟实际资源的情况。

Answer 1

可以按如下方式访问和更改状态：

athena_backends['us-east-2'].executions.get(query_id).status

示例代码片段

from moto.athena import athena_backends
query = "SELECT stuff"
location = "s3://bucket-name/prefix/"
database = "database"
# Start Query
exex_id = self.client.start_query_execution(
   QueryString=query,
   QueryExecutionContext={"Database": database},
   ResultConfiguration={"OutputLocation": location},
)["QueryExecutionId"]
athena_backends['us-west-2'].executions.get(exex_id).status = "CANCELLED"