Mock AWSwrangler for unittesting

user_10131384 0
Viewed 513

As there is no support for AWSwrangler by moto i am stuck here and don't know how to mock.

I am trying to unittest my lambda code which run athena query using AWSwrangler.

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())

I referenced AWSwrangler github issue and while trying some of the test provided in link it's hitting AWS service instead of running locally.

amazon-web-services unit-testing python
1 Answers

Here is an example implementation for this function using moto and pytest.

First I would correct your function according to awswrangler required parameters in its current version (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

The in a test/conftest.py fil I would declare the necessary mocked objects:

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/"}
        )

And, then I would define a test of the function in a separate test/athena_test.py file. This is using mocker to mock awswrangler response to the query but you could use advanced testing using the mock objects created in the conftest.py file :

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

Resources:

user_143783140
Related