用于单元测试的模拟 AWSwrangler答案

【问题标题】：Mock AWSwrangler for unittesting用于单元测试的模拟 AWSwrangler
【发布时间】：2022-08-02 21:21:17
【问题描述】：

由于 moto 不支持 AWSwrangler 我被困在这里并且不知道如何模拟。

我正在尝试对我的 lambda 代码进行单元测试，该代码使用 AWSwrangler 运行 athena 查询。

import awswrangler as wr
import boto3

def athena_query(dbtable, contact_id, athena_output, session):
    
    query = \"\"\"
    SELECT
        *
    FROM
        :dbtable;
    WHERE 
    contactid=:contactid;
    \"\"\"

    output = wr.athena.read_sql_query(
        query, 
        params = {
            \"contactid\": f\"\'{contact_id}\'\", 
            \"dbtable\": f\"{dbtable}\"
        }, 
        s3_output = athena_output,
        boto3_session = session
    )
    results = output.head().loc[0]
    
    return results

response = athena_query(\"table_name\", \"123\", \"s3://bucket\", boto3.session.Session())

我引用了AWSwrangler github issue，并在尝试链接中提供的一些测试时，它正在访问 AWS 服务，而不是在本地运行。

请添加完整的代码（包括导入了哪些包）并改进格式。
@slideshowp2 我按照您的建议进行了一些更改。

标签： python amazon-web-services unit-testing

【解决方案1】：

这是使用 moto 和 pytest 实现此功能的示例。

首先，我将根据当前版本（2.16.1）中的awswrangler 所需参数更正您的功能。

import awswrangler as wr
import boto3

def athena_query(database, dbtable, contact_id, athena_output, session):
    
    query = """
    SELECT
        *
    FROM
        :dbtable;
    WHERE 
    contactid=:contactid;
    """

    output = wr.athena.read_sql_query(
        query, 
        database,
        params = {
            "contactid": f"'{contact_id}'", 
            "dbtable": f"{dbtable}"
        }, 
        s3_output = athena_output,
        boto3_session = session
    )
    results = output.head().loc[0]
    
    return results

在test/conftest.py 文件中，我将声明必要的模拟对象：

import pytest
import moto


TEST_BUCKET_NAME = "my_bucket"
REGION = "us-east-1"
DATABASE_NAME = "test_db"
TABLE_NAME = "test_table"
TABLE_DDL = f"""CREATE EXTERNAL TABLE IF NOT EXISTS
{DATABASE_NAME}.{TABLE_NAME} (
  a string,
  b string,
  contactid string
) ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'
WITH SERDEPROPERTIES (
  'separatorChar' = ',',
  'quoteChar' = '\"',
  'escapeChar' = '\\'
)
STORED AS TEXTFILE
LOCATION 's3://{TEST_BUCKET_NAME}/input/';"""


@pytest.fixture
def aws_credentials():
    """Mocked AWS Credentials for moto."""
    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
    os.environ["AWS_SECURITY_TOKEN"] = "testing"
    os.environ["AWS_SESSION_TOKEN"] = "testing"


@pytest.fixture
def s3_client(aws_credentials):
    with moto.mock_s3():
        conn = boto3.client("s3", region_name=REGION)
        yield conn

@pytest.fixture
def s3_client(aws_credentials):
    with moto.mock_s3():
        conn = boto3.client("s3", region_name=REGION)
        yield conn

@pytest.fixture
def athena_client(aws_credentials):
    with moto.athena.mock_athena():
        conn = boto3.client("athena", region_name=REGION)
        yield conn

@pytest.fixture
def s3_bucket(s3_client):
    s3_client.create_bucket(
        Bucket=TEST_BUCKET_NAME,
        CreateBucketConfiguration={
            'LocationConstraint': 'eu-west-1'
        }
    )
    yield boto3.resource('s3').Bucket(TEST_BUCKET_NAME)

@pytest.fixture
def athena_table(athena_client, s3_bucket):
    # create database
    _ = athena_client.start_query_execution(
        QueryString=f"create database {DATABASE_NAME}",
        ResultConfiguration={"OutputLocation": "s3://{TEST_BUCKET_NAME}/queries/"}
    )
    # create table 
    _ = athena_client.start_query_execution(
            QueryString=TABLE_DDL,
            ResultConfiguration={"OutputLocation": "s3://{TEST_BUCKET_NAME}/queries/"}
        )

然后，我将在单独的 test/athena_test.py 文件中定义函数测试。这是使用 mocker 来模拟 awswrangler 对查询的响应，但您可以使用在 conftest.py 文件中创建的模拟对象进行高级测试：

from conftest import TEST_BUCKET_NAME, DATABASE_NAME, TABLE_NAME
# import your function to test here

def test_athena_query(s3_bucket, athena_table, mocker):

    def mock_response(*args, **kwargs):
        return pd.DataFrame.from_dict({"a": [1, 2], "b": [3, 4], "contactid": [123, 123]})

    # mocking
    mock_wr_call = mocker.patch('wr.athena.read_sql_query')
    mock_wr_call.side_effect = mock_response

    response = athena_query(DATABASE_NAME, TABLE_NAME, "123", f"s3://{TEST_BUCKET_NAME}/queries/", boto3.session.Session())

    assert response.shape[0] == 2

资源：

https://aws-data-wrangler.readthedocs.io/en/stable/stubs/awswrangler.athena.read_sql_query.html

【讨论】：