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.
51 lines
1.6 KiB
51 lines
1.6 KiB
2 months ago
|
import unittest
|
||
|
|
from unittest.mock import patch
|
||
|
|
import jieba
|
||
|
|
import wordcloud
|
||
|
|
import pandas as pd
|
||
|
|
|
||
|
|
class TestWordCloudGeneration(unittest.TestCase):
|
||
|
|
|
||
|
|
@patch('pandas.read_excel')
|
||
|
|
@patch('wordcloud.WordCloud.to_file')
|
||
|
|
def test_wordcloud_generation(self, mock_to_file, mock_read_excel):
|
||
|
|
# 模拟读取的 Excel 数据
|
||
|
|
mock_df = pd.DataFrame({
|
||
|
|
'AI': ['人工智能在发展', 'AI改变世界'],
|
||
|
|
'Machine learning': ['机器学习是AI的基础', '深度学习是其中之一']
|
||
|
|
})
|
||
|
|
mock_read_excel.return_value = mock_df
|
||
|
|
|
||
|
|
# 引入你需要测试的代码
|
||
|
|
df = pd.read_excel('top_ai_danmakus.xlsx')
|
||
|
|
|
||
|
|
# 将Excel所有列的数据拼接成一个长字符串
|
||
|
|
book = ''
|
||
|
|
for column in df.columns:
|
||
|
|
book += ' '.join(df[column].astype(str)) + ' '
|
||
|
|
|
||
|
|
# 使用jieba进行分词
|
||
|
|
book_list = jieba.lcut(book)
|
||
|
|
book_str = ' '.join(book_list)
|
||
|
|
|
||
|
|
# 创建词云对象
|
||
|
|
wc = wordcloud.WordCloud(
|
||
|
|
width=500,
|
||
|
|
height=500,
|
||
|
|
background_color='white',
|
||
|
|
stopwords={'main', 'Taipei', 'afraid', 'aiden', 'Britain'}, # 可添加不需要的停用词
|
||
|
|
font_path='msyh.ttc' # 设置中文字体路径
|
||
|
|
)
|
||
|
|
|
||
|
|
# 生成词云
|
||
|
|
wc.generate(book_str)
|
||
|
|
|
||
|
|
# 检查是否调用了保存词云的函数
|
||
|
|
wc.to_file('词云图.png')
|
||
|
|
|
||
|
|
# 验证 to_file 函数是否被调用
|
||
|
|
mock_to_file.assert_called_once_with('词云图.png')
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|