Justin S
Justin S

Reputation: 1459

Using AWS Moto with Python mock to write unit tests

I am working with a codebase that doesn't have any tests setup. I am trying to add some tests and currently have the below test class test_main.py

from unittest import TestCase
from moto import mock_aws
import boto3

from main import collect_emr


class Test(TestCase):

    def setUp(self):
        pass

    @mock_aws
    def test_collect_emr(self):
        emr_client = boto3.client('emr')

        emr_client.run_job_flow(
            Instances={
                "InstanceCount": 3,
                "KeepJobFlowAliveWhenNoSteps": True,
                "MasterInstanceType": "c3.medium",
                "Placement": {"AvailabilityZone": "us-east-1a"},
                "SlaveInstanceType": "c3.xlarge",
            },
            JobFlowRole="EMR_EC2_DefaultRole",
            LogUri="s3://mybucket/log",
            Name="cluster",
            ServiceRole="EMR_DefaultRole",
            VisibleToAllUsers=True,
            Tags=[
                {
                    'Key': 'string',
                    'Value': 'string'
                },
            ]
        )

        expected_outcome = [{'arn': 'arn:aws:elasticmapreduce:ap-southeast-2:123456789012:cluster/j-S78L4S34G4Q7U', 'tags': [{'key': 'string', 'value': 'string'}]}]
        actual_outcome = list(collect_emr('us-east-1', 'arn:aws:iam::123456789012:role/role_name'))

        self.assertEqual(expected_outcome, actual_outcome)

then the method I am testing from main.py:

def collect_emr(region_name: str, role_arn: str) -> Generator[ResourceInfo, Any, None]:
    creds = assume_role(role_arn)
    client = boto3.client(
        "emr",
        aws_access_key_id=creds["AccessKeyId"],
        aws_secret_access_key=creds["SecretAccessKey"],
        aws_session_token=creds["SessionToken"],
        region_name=region_name,
        config=config,
    )

    cluster_ids = []
    for page in client.get_paginator("list_clusters").paginate():
        for cluster in page['Clusters']:
            if 'Id' not in cluster:
                continue
            else:
                cluster_id = cluster['Id']
                cluster_ids.append(cluster_id)
    
    if cluster_ids:
        for cluster_id in cluster_ids:
            describe_cluster = client.describe_cluster(ClusterId = cluster_id)
            if 'Cluster' not in describe_cluster:
                continue
            yield {
                "arn": describe_cluster['Cluster']['ClusterArn'],
                "tags": list(
                    [{"key": t["Key"], "value": t["Value"]} for t in describe_cluster['Cluster']["Tags"]]
                ),
            }

and my test output when I run them:

Ran 1 test in 0.852s

FAILED (failures=1)


[] != [{'arn': 'arn:aws:elasticmapreduce:ap-southeast-2:123456789012:cluster/j-S78L4S34G4Q7U',
  'tags': [{'key': 'string', 'value': 'string'}]}]

I understand that the client = inside the collect_emr(... method in main.py is creating a new boto3 client which doesn't know about the resources created by emr_client created in my test method test_collect_emr(self): but im unsure how to go about resolving this? In the past I would write methods that would always take the client as a parameter eg:

def collect_emr(boto_client):
    ...

and then pass in the boto client I created in my test case. I am unable to do that here as the code base has stuff all over the place and refactoring the method to do things this way would break a lot of stuff. Is it possible to combine python mock and aws moto in the same method or would I need to use either or? Any help would be greatly appreciated as im a bit stuck on the best way to approach this

Upvotes: 0

Views: 201

Answers (0)

Related Questions