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

Question

我有一个文件seed_dynamodb.py，其代码如下。我想用mock patch编写一个单元测试。我能够成功地修补boto3。现在我也需要给sys.argv打补丁。我已经尝试了下面的测试代码，但它给出了一个IndexError。

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

I am getting below issue:

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或需要修补任何东西。

Answer 3

我试着修改你的代码来执行和解决IndexError: list index out of range.
为了在我的系统中尝试测试代码，我对你的文件做了一些改变。所有的更改都通过一些注释突出显示。
基本上我已经改变了你的导入位置:from module import seed_dynamodb．此外(没有在评论中突出显示)我已经更改了您的print(f"{..}")说明，因为它们不正确。

========== 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 = {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]      # <-------- your code (you will set index to 1 in your 
#           definitive code)
env = sys.argv[0]       # <-------- my code
#region = sys.argv[2]   # <-------- your code (you will set index to 2 in your
#           definitive code)
region = sys.argv[1]    # <-------- my code
l = len(sys.argv)
print(f'env={env}, region={region}, l={l}') # <--- change a bit the print(f"...")
main(env, region)

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

#from module import seed_dynamodb    # <-------- your code
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):
from module import seed_dynamodb    # <-------- my code
mock_boto = args[0]
#mock_dynamo = args[0]  # <-------- your code
mock_dynamo = args[1]   # <-------- my code
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)

这是在我的系统中执行测试的输出:

test_seed_dynamodb (__main__.SeedDynamoDBTest) ... ok
----------------------------------------------------------------------
Ran 1 test in 0.002s
OK
env=pr, region=us-east-1, l=2
table_name = syn-fcad-nielsen-pr-time

我希望这个答案对你有用。