Python Mock 测试

import os
import os.path

def rm_file(filename):
if os.path.isfile(filename):
os.remove(filename)

# 当被测函数在另外一个模块里时，需要导入

from mymodule import rm_file

import os.path
import tempfile
import unittest

class RmTestCase(unittest.TestCase): # 创建临时文件
tmpfilepath = os.path.join(tempfile.gettempdir(), "tmp-testfile")

    # 向临时文件中写入测试数据
    def setUp(self):
        with open(self.tmpfilepath, "wb") as f:
            f.write("Delete me!")

    def test_rm(self):
        # 调用目标删除函数
        rm_file(self.tmpfilepath)
        # 判断文件是否还存在
        self.assertFalse(os.path.isfile(self.tmpfilepath), "Failed to remove the file.")

# 当被测函数在另外一个模块里时，需要导入

from mymodule import rm_file

import mock
import unittest

class TestRemoveFile(unittest.TestCase):
@patch('os.remove') # 模拟对象 os.remove
@patch('os.path') # 模拟对象 os.path
def test_rm(self, mock_path, mock_remove): # mock 对象 mock_remove 会替换 os.remove 函数 # 模拟对象 mock_path.isfile 返回值为 True, 表示文件存在
mock_path.isfile.return_value = True # 被测函数调用
rm_file('dummy.txt')

        # 断言参数被 os.path.isfile(), os.remove() 调用，且只调用一次
        mock_path.isfile.assert_called_once_with('dummy.txt')
        mock_remove.assert_called_once_with('dummy.txt')

    @patch('os.remove')
    @patch('os.path.isfile')  # 模拟对象 os.path.isfile， 跟模拟 os.path 差不多
    def test_rm_file_does_not_exist(self, mock_isfile, mock_remove): # 这里参数的顺序不能错，按照装饰器 @patch 从下到上的顺序在这里从左到右
        # 模拟当文件不存在时
        mock_isfile.return_value = False
        # 被测函数调用
        rm_file('dummy.txt')

        # 只调用了 os.path.isfile() 一次, 但是没有调用 os.remove()，因为文件不存在
        mock_isfile.assert_called_once_with('dummy.txt')
        mock_remove.assert_not_called()

class unittest.mock.Mock(spec=None, side_effect=None, return_value=DEFAULT, wraps=None, name=None, spec_set=None, unsafe=False, \*\*kwargs)

# 假设有另外一个函数

def mock_func(): # 做一些操作，然后返回一个值
return 30

# 1. 创建 mock 对象可以这样

mock_obj = mock.Mock(return_value=20,side_effect=mock_func, name='mock_obj')

# 2. 或者这样：

mock_obj = mock.Mock()
mock_obj.return_value = 20
mock_obj.side_effect = mock_func

# mock_func 和 return_value 同时存在，所以模拟替换之后，返回的值是 30

unittest.mock.patch(target, new=DEFAULT, spec=None, create=False, spec_set=None, autospec=None, new_callable=None, \*\*kwargs)

def cleanup_backups(count, backup_path=default_backup_dir):
"""
cleanup outdated backup files
删除指定路径下指定数量的备份文件
"""
if count < 0:
return
backups = glob.glob("{}/\*.gz".format(backup_path))
backups.sort(reverse=True) # 遍历最老的 count 个文件
for f in backups[count:]:
md5sum_file = f.replace(".gz", "") + ".md5sum"
os.remove(f) # 如果备份文件有对应的 md5sum 文件，也进行删除
if os.path.exists(md5sum_file):
os.remove(md5sum_file)
else: # 没有就不删除
logging.error("File {} not found, skipping deletion".format(md5sum_file))
continue # 输出 log 信息
logging.info("Deleting expired backup file {}".format(f))
logging.info("Deleting expired backup file {}".format(md5sum_file))

class TestCleanUpBackups(unittest.TestCase): # 假设的备份文件路径，不需要提供真实路径
backup_path = "/path/to/backup/pg-data" # 假如有以下备份文件，供模拟对象使用，文件名表示时间戳
backup_files = [
"/path/to/backup/pg-data/20231020.gz",
"/path/to/backup/pg-data/20231019.gz",
"/path/to/backup/pg-data/20231018.gz",
"/path/to/backup/pg-data/20231017.gz",
]

    @patch("glob.glob")
    @patch("os.path.exists")
    @patch("os.remove")
    @patch("logging.error")
    @patch("logging.info")
    def test_cleanup_backups_01(
        self, mock_info, mock_error, mock_remove, mock_exists, mock_glob
    ):
        mock_glob.return_value = backup_files
        # mock_exists 对应的是 os.path.exists
        # 因为在循环了只被调用两次，所有其结果有两次值，[True, True] 代表两次判断都是文件存在. [False, True]代表有个文件不存在。
        # 所以这里可以控制产生多个测试 case
        # 测试有一个文件不存在对应的 .md5sum 文件，则跳过删除
        mock_exists.side_effect = [False, True]
        # 只删除最老的 2 份文件
        cleanup_backups(2, self.backup_path)

        mock_glob.assert_called_with("{}/*.gz".format(self.backup_path))
        # 跳过删除
        mock_error.assert_has_calls(
            [
                call(
                    "File /path/to/backup/pg-data/20231018.md5sum not found, skipping deletion"
                )
            ]
        )
        # 所以实际只删除一个备份文件 20231017
        mock_remove.assert_has_calls(
            [
                call("/path/to/backup/pg-data/20231017.gz"),
            ]
        )
        # 最后分别调用了以下 logging.info
        mock_info.assert_has_calls(
            [
                call(
                    "Deleting expired backup file /path/to/backup/pg-data/20231017.gz"
                ),
                call(
                    "Deleting expired backup file /path/to/backup/pg-data/20231017.md5sum"
                ),
            ]
        )

import unittest
from unittest.mock import patch
import requests

def get_data_from_api():
response = requests.get('https://api.example.com/data')
return response.json()

class TestGetDataFromApi(unittest.TestCase):
@patch('requests.get')
def test_get_data_from_api(self, mock_get): # 模拟请求返回的 response
mock_response = mock_get.return_value # 模拟返回状态码为 200
mock_response.status_code = 200
mock_response.json.return_value = {'data': 'test'}

        result = get_data_from_api()
        assert result['data'] == 'test'

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
import unittest
from unittest.mock import patch

from models import User # 假设有一个 User 模型

# 假设需要根据 id 查询用户信息

def get_user_by_id(user_id):
engine = create_engine('sqlite:///test.db')
Session = sessionmaker(bind=engine)
session = Session()
user = session.query(User).filter_by(id=user_id).first()
return user
class TestQueryFromDB(unittest.TestCase):
def test_get_user_by_id(self): # 模拟一个 models.Session 对象。 # 并且这里使用 with 语法将 mock 对象用作上下文，但是这样只适合有一个模拟对象的情况
with patch('models.Session') as mock_session: # 模拟 session 返回的一个实例
mock_instance = mock_session.return_value # 模拟实例的 query, filter_by, first 方法的返回值，模拟查询结果
mock_instance.query.return_value.filter_by.return_value.first.return_value = User(id=1, name='Alice')

            result = get_user_by_id(1)
            assert result.id == 1
            assert result.name == 'Alice'

import requests
from unittest.mock import patch
import docker

# 获取指定容器

def get_container(name):
client = docker.DockerClient(base_url="unix://var/run/docker.sock")
try:
container = client.containers.get(name)
except docker.errors.NotFound as e:
logging.error("container {} not found, error: {}".format(name, e))
return None
return container

# 在指定容器中运行命令

def run_cmd_in_container(command, container, user="root"):
if container is None:
return False, None

    exit_code, output = container.exec_run(command, user=user, privileged=True)
    if exit_code != 0:
        logging.error(
            "Failed to run {} in container , exit code: {}, output: {}".format(
                command, exit_code, output
            )
        )
        return False, None
    return True, output

class TestDocker(unittest.TestCase):
not_found_exception = docker.errors.NotFound("Container not found") # 创建 mock 对象，并且模拟没有找到容器的错误 case
docker_client_mock = Mock()
docker_client_mock.containers.get.side_effect = not_found_exception

    @patch("docker.DockerClient", return_value=docker_client_mock)
    @patch("logging.error")
    def test_get_container(self, mock_error, docker_mock):
        container = get_container("sds-postgres")

        self.assertIsNone(container)
        docker_mock.assert_called_once()
        mock_error.assert_called_once()
        mock_error.assert_has_calls(
            [call("container sds-postgres not found, error: Container not found")]
        )
    # 最后清除 mock 对象资源：reset_mock()
    docker_client_mock.reset_mock()

    @patch("docker.DockerClient")
    def test_run_cmd_in_container(self, mock_docker_client):
        # 创建 Mock 对象， 模拟一个 docker client
        mock_client = Mock()
        mock_docker_client.return_value = mock_client
        # 创建 Mock 对象， 模拟一个 docker 容器
        mock_container = Mock()
        mock_client.containers.get.return_value = mock_container

        mock_container.exec_run.return_value = (0, b"Success")
        ret, output = run_cmd_in_container("psql", mock_container)

        mock_container.exec_run.assert_called_with("psql", user="root", privileged=True)
        self.assertEqual(output, b"Success")
        self.assertTrue(ret)