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