如何在 python 单元测试中使用模拟补丁正确修补 sys.argv

Question

我有文件 seed_dynamodb.py，其代码如下。我想使用模拟补丁对此进行单元测试。我成功地修补了 boto3。现在我还需要修补 sys.argv。我尝试了下面的测试代码，但它给出了错误

==========seed_dynamodb.py============

import sys
import boto3
def main(env,region):

    dynamodb_client= boto3.client('dynamodb')
    timestamp = '1234567'
    table_name = 'syn-fcad-nielsen-' + env + '-time'
    print(f'{table_name=}')
    if env == 'uat':
     timestamp = 1234567
    if env == 'prod':
     timestamp = 1234567
    response = dynamodb_client.get_item(TableName=table_name, 
                                        Key={'BaseTime':{'S':'Timestamp'}})
    if 'Item' in response:
        print("Item exists in Dynamo DB table")
        timestamp = response['Items']['Value']['N']
    else: 
        response = dynamodb_client.put_item(TableName=table_name, 
                    Item={
                    'BaseTime':{'S':'Timestamp'},
                    'Value': {'N': timestamp}
                    })

env = sys.argv[1]
region = sys.argv[2]
l = len(sys.argv)
print(f'{env=}{region=}{l=}')
main(env,region)

=======================test_dynamo.py=========

 from module import seed_dynamodb
    import unittest
    from unittest import mock
    from unittest.mock import patch
    import boto3
    import sys
    @mock.patch("module.seed_dynamodb.boto3.client")
    class SeedDynamoDBTest(unittest.TestCase):
        @patch.object(boto3, "client")
        @patch.object(sys, 'argv', ["pr", "us-east-1"])
        def test_seed_dynamodb(self, *args):
            mock_boto = args[0]
            mock_dynamo = args[0]
            mock_dynamo.get_item.return_value = {"Item": {"Value": {"N": 1678230539}}}
            mock_dynamo.put_item.return_value = {
                "Item": {"BaseTime": {"S": "Timestamp"}, "Value": {"N": 1678230539}}
            }
            seed_dynamodb.dynamodb_client = mock_dynamo
            self.assertIsNotNone(mock_boto.client.get_item.return_value)
           # seed_dynamodb.main("pr-173", "us-east-1")
    
    
    if __name__ == "__main__":
        unittest.main(verbosity=2)

我遇到以下问题

    env = sys.argv[1]
IndexError: list index out of range

你能帮我解决这个问题或编写测试用例来修补 sys.argv

Answer 1

让我们尝试直接修补

sys.argv

而不是使用

patch.object

这是您更新的代码：

from module import seed_dynamodb
import unittest
from unittest import mock
from unittest.mock import patch
import boto3
import sys

@mock.patch("module.seed_dynamodb.boto3.client")
class SeedDynamoDBTest(unittest.TestCase):
    @patch.object(boto3, "client")
    def test_seed_dynamodb(self, mock_boto):
        mock_dynamo = mock_boto.return_value
        mock_dynamo.get_item.return_value = {"Item": {"Value": {"N": 1678230539}}}
        mock_dynamo.put_item.return_value = {
            "Item": {"BaseTime": {"S": "Timestamp"}, "Value": {"N": 1678230539}}
        }
        with patch("sys.argv", ["script_name", "pr", "us-east-1"]):
            seed_dynamodb.main(sys.argv[1], sys.argv[2])
        self.assertIsNotNone(mock_dynamo.client.get_item.return_value)

if __name__ == "__main__":
    unittest.main(verbosity=2)

这应该允许您测试

seed_dynamodb.main

的功能

Answer 2

首先不要让你的代码如此直接地依赖于

sys.argv

。而不是

env = sys.argv[1]
region = sys.argv[2]
l = len(sys.argv)
print(f'{env=}{region=}{l=}')
main(env,region)

做类似的事情

def parse_args(argv=None):
    p = argparse.ArgumentParser()
    p.add_argument('env')
    p.add_argument('region')
    return p.parse_args(argv)

if __name__ == '__main__':
    args = parse_args()
    main(args.env, args.region)

现在您可以使用您喜欢的任何列表来测试

parse_args

，而不是依赖于

sys.argv

或需要修补任何东西。