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