python - 模拟 AWSwrangler 进行单元测试

Question

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

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

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 服务而不是在本地运行。

Answer 1

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

首先，我会根据 awswrangler 当前版本 (2.16.1) 中所需的参数来更正您的函数。

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