ADD 单元测试代码

test_word_cloud.py

@@ -0,0 +1,58 @@
+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()