note
- 使用
unittest框架进行单元测试是Python标准库的一部分，提供了编写测试用例、测试套件以及运行测试的能力。 - 测试用例是继承自
unittest.TestCase的类。在这个类中，你可以定义一系列的方法来测试不同的行为。每个测试方法都应该以test开头。
一、单元测试unittest框架
1. 直接来看一个简单的测试用例
- 使用
unittest框架进行单元测试是Python标准库的一部分，提供了编写测试用例、测试套件以及运行测试的能力。 - 测试用例是继承自
unittest.TestCase的类。在这个类中，你可以定义一系列的方法来测试不同的行为。每个测试方法都应该以test开头。下面代码是一个简单的测试用例
# test-单元测试
import unittest
# 子类必须继承unittest.TestCase类
class TestMethod(unittest.TestCase):
# 每个测试方法都需要以test开头
def test_upper(self):
# 检查是否相等
self.assertEqual('foo'.upper(), 'FOO')
def test_isupper(self):
# 验证是否为真/假
self.assertTrue('FOO'.isupper())
self.assertFalse('foo'.isupper())
def test_split(self):
s = "happy new year"
self.assertEqual(s.split(), ['happy', 'new', 'year'])
with self.assertRaises(TypeError):
s.split(3)
if __name__ == '__main__':
unittest.main()
2. 相关断言方法的使用：
- 使用
assertEqual(a, b)来检查a和b是否相等。 - 使用
assertTrue(x)和assertFalse(x)来验证x是否为真或假。 - 使用
assertRaises(Error, func, *args, **kwargs)来验证是否抛出了期望的异常。比如上面的代码例子，s.split()方法的参数应该是一个字符串，然后传入了3数字所以预期该调用会抛出TypeError异常
3. 运行测试
可以通过以下2种方式运行测试：
- a. 命令行运行
如果你的测试代码保存在test.py文件中，可以直接在命令行中运行：
python -m unittest test
这会自动发现所有以test开头的方法并运行它们。
- b. 在代码中运行
如果你想在代码中直接运行测试，可以在文件的最后加上：
if __name__ == '__main__':
unittest.main()
4. 测试套件和测试运行器
对于更复杂的测试需求，可以使用unittest.TestSuite来组合多个测试用例，然后用unittest.TextTestRunner来运行这些测试。
def suite():
suite = unittest.TestSuite()
suite.addTest(TestStringMethods('test_upper'))
suite.addTest(TestStringMethods('test_isupper'))
return suite
if __name__ == '__main__':
runner = unittest.TextTestRunner()
runner.run(suite())
unittest框架提供了丰富的断言方法和测试组织方式，帮助开发者构建可靠和可维护的测试套件。通过这种方式，可以有效地进行单元测试，确保代码的质量和稳定性。
二、一个完整的例子
这里引用大模型数据处理库data_juicer的一个栗子，PerplexityFilter是根据计算微调数据的困惑度进行筛选sft数据：
import unittest
from datasets import Dataset
from data_juicer.ops.filter.perplexity_filter import PerplexityFilter
from data_juicer.utils.constant import Fields
class PerplexityFilterTest(unittest.TestCase):
def _run_perplexity_filter(self, dataset: Dataset, target_list, op):
if Fields.stats not in dataset.features:
# TODO:
# this is a temp solution,
# only add stats when calling filter op
dataset = dataset.add_column(name=Fields.stats,
column=[{
}] * dataset.num_rows)
dataset = dataset.map(op.compute_stats)
dataset = dataset.filter(op.process)
dataset = dataset.select_columns(column_names=['text'])
res_list = dataset.to_list()
self.assertEqual(res_list, target_list)
def test_en_case(self):
ds_list = [{
'text': "Today is Sunday and it's a happy day!"
}, {
'text':
"Today is Sund Sund Sund Sund Sunda and it's a happy day!"
}, {
'text': 'a v s e c s f e f g a qkc'
}, {
'text': '，。、„”“«»１」「《》´∶：？！（）；–—．～’…━〈〉【】％►'
}, {
'text': 'Do you need a cup of coffee?'
}, {
'text': 'emoji表情测试下😊，😸31231'
}]
tgt_list = [{
'text': "Today is Sunday and it's a happy day!"
}, {
'text': 'Do you need a cup of coffee?'
}]
dataset = Dataset.from_list(ds_list)
op = PerplexityFilter(lang='en', max_ppl=900)
self._run_perplexity_filter(dataset, tgt_list, op)
if __name__ == '__main__':
unittest.main()
Reference
[1] python中unittest库 python的unittest框架
[2] https://docs.python.org/zh-cn/3.11/library/unittest.html
