Блог

Я использую 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, и я не могу найти способ явно изменить его. В идеале это могло бы произойти где-нибудь в тесте, чтобы вручную изменить состояние запроса, чтобы смоделировать, как это будет с реальными ресурсами.

Доступ к состоянию и его изменение можно получить следующим образом:

 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"
 

Мне кажется, что moto возвращает только QUEUED для start_query_execution , вы можете взглянуть на исходный код здесь.

Используется другой подход from unittest import mock , и тогда вы можете сделать что-то вроде:

 cls.CLIENT = mock.Mock()
cls.CLIENT.start_query_execution.side_effect = [
    'QUEUED',
    'SUCCEEDED'
]

 

Итак, при первом cls.CLIENT.start_query_execution(..) вызове он вернет сообщение о том, что запрос находится в очереди, но во второй раз вернет сообщение об успешном выполнении, и тогда вы сможете протестировать оба пути выполнения.

Кроме того, с moto его помощью вы не сможете протестировать все случаи, потому что, кроме статуса в очереди, вы можете установить статус запроса только на CANCELLED , как вы можете видеть здесь .