You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
59 lines
1.8 KiB
59 lines
1.8 KiB
import unittest
|
|
from unittest.mock import patch, mock_open
|
|
import jieba
|
|
import wordcloud
|
|
|
|
class TestWordCloudGeneration(unittest.TestCase):
|
|
|
|
@patch('builtins.open', new_callable=mock_open, read_data='AI技术是未来,GPT非常强大')
|
|
def test_file_reading(self, mock_file):
|
|
"""
|
|
测试文件读取功能
|
|
"""
|
|
|
|
with open('AI统计弹幕.txt', encoding='utf-8') as f:
|
|
content = f.read()
|
|
mock_file.assert_called_once_with('AI统计弹幕.txt', encoding='utf-8')
|
|
self.assertEqual(content, 'AI技术是未来,GPT非常强大')
|
|
|
|
def test_jieba_cut(self):
|
|
"""
|
|
测试 jieba 分词功能
|
|
"""
|
|
test_text = 'AI技术是未来,GPT非常强大'
|
|
expected_result = ['AI', '技术', '是', '未来', ',', 'GPT', '非常', '强大']
|
|
|
|
# 调用 jieba.lcut() 进行分词
|
|
result = jieba.lcut(test_text)
|
|
|
|
self.assertEqual(result, expected_result)
|
|
|
|
|
|
@patch('wordcloud.WordCloud.to_file')
|
|
def test_wordcloud_stopwords(self, mock_to_file):
|
|
"""
|
|
测试词云是否不包含停用词
|
|
"""
|
|
test_text = 'AI技术是未来,GPT非常强大'
|
|
txt_list = jieba.lcut(test_text)
|
|
string = ' '.join(txt_list)
|
|
|
|
wc = wordcloud.WordCloud(
|
|
stopwords={'是', '未来'} # 停用词
|
|
)
|
|
|
|
wc.generate(string)
|
|
wc.to_file('out.png')
|
|
|
|
# 检查生成的词云中是否包含停用词
|
|
self.assertIn('AI', wc.words_)
|
|
self.assertIn('GPT', wc.words_)
|
|
self.assertNotIn('是', wc.words_)
|
|
self.assertNotIn('未来', wc.words_)
|
|
|
|
# 验证文件保存调用
|
|
mock_to_file.assert_called_once_with('out.png')
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|