ADD file via upload

tests/test_visualization.py

@@ -0,0 +1,137 @@
+"""
+test_visualization.py - 单元测试可视化模块的功能。
+"""
+import unittest
+from unittest.mock import patch
+import pandas as pd
+from visualization import (
+    read_danmakus,
+    filter_relevant_danmakus,
+    preprocess_danmakus,
+    extract_words,
+    remove_stopwords,
+    calculate_word_frequency,
+    generate_wordcloud
+)
+
+class TestVisualization(unittest.TestCase):
+    """
+    测试可视化模块的单元测试类。
+    """
+
+    @patch('pandas.read_excel')
+    def test_read_danmakus(self, mock_read_excel):
+        """
+        测试从 Excel 文件中读取弹幕内容的功能。
+        """
+        # 模拟 Excel 数据
+        mock_data = pd.DataFrame({'弹幕内容': ['弹幕1', '弹幕2', '弹幕3']})
+        mock_read_excel.return_value = mock_data
+
+        # 调用读取函数
+        result = read_danmakus("dummy_filename.xlsx")
+        # 期望返回弹幕内容列表
+        self.assertEqual(result, ['弹幕1', '弹幕2', '弹幕3'])
+
+    def test_filter_relevant_danmakus(self):
+        """
+        测试筛选与给定关键词相关的弹幕内容。
+        """
+        danmakus = ['这是AI技术的未来', '普通弹幕', '智能应用']
+        keywords = ['AI技术', '智能']
+
+        # 调用筛选函数
+        result = filter_relevant_danmakus(danmakus, keywords)
+        # 期望返回包含关键词的弹幕
+        self.assertEqual(result, ['这是AI技术的未来', '智能应用'])
+
+    def test_filter_relevant_danmakus_no_match(self):
+        """
+        测试没有匹配关键词时的筛选结果。
+        """
+        danmakus = ['普通弹幕1', '普通弹幕2']
+        keywords = ['AI技术', '智能']
+
+        # 调用筛选函数
+        result = filter_relevant_danmakus(danmakus, keywords)
+        # 期望返回空列表
+        self.assertEqual(result, [])
+
+    def test_preprocess_danmakus(self):
+        """
+        测试预处理弹幕内容，去除特殊字符和多余空格。
+        """
+        danmakus = ['这是一个测试！', '普通  弹幕  ']
+
+        # 调用预处理函数
+        result = preprocess_danmakus(danmakus)
+        # 期望返回处理后的弹幕
+        self.assertEqual(result, ['这是一个测试', '普通 弹幕'])
+
+    def test_extract_words(self):
+        """
+        测试使用 Jieba 分词库对弹幕内容进行分词。
+        """
+        danmakus = ['这是一个测试']
+
+        # 调用分词函数
+        result = extract_words(danmakus)
+        # 期望返回分词结果
+        self.assertIn('这是', result)
+        self.assertIn('一个', result)
+        self.assertIn('测试', result)
+
+    def test_remove_stopwords(self):
+        """
+        测试去除停用词和单个字符。
+        """
+        words = ['的', '是', '测试', 'AI', '智能']
+        stopwords = set(['的', '是'])
+
+        # 调用去除停用词函数
+        result = remove_stopwords(words, stopwords)
+        # 期望返回去除停用词后的词语列表
+        self.assertEqual(result, ['测试', 'AI', '智能'])
+
+    def test_calculate_word_frequency(self):
+        """
+        测试计算词语的频率。
+        """
+        words = ['测试', 'AI', '测试']
+
+        # 调用计算词频函数
+        result = calculate_word_frequency(words)
+        # 期望返回正确的频率计数器
+        self.assertEqual(result['测试'], 2)
+        self.assertEqual(result['AI'], 1)
+
+    @patch('visualization.WordCloud')
+    @patch('visualization.plt')
+    def test_generate_wordcloud(self, mock_plt, mock_wordcloud):
+        """
+        测试生成词云图并将其保存为 PNG 文件的功能。
+        """
+        word_freq = {'测试': 10, 'AI': 5}
+
+        # 调用生成词云函数
+        generate_wordcloud(word_freq)
+
+        # 确保 WordCloud 的生成函数被调用
+        mock_wordcloud.assert_called_once()
+        # 确保 plt 的 savefig 被调用
+        mock_plt.savefig.assert_called_once_with("wordcloud.png", bbox_inches='tight', dpi=300)
+
+    def test_generate_wordcloud_no_words(self):
+        """
+        测试生成词云图时没有词语的情况。
+        """
+        word_freq = {}
+
+        # 调用生成词云函数
+        try:
+            generate_wordcloud(word_freq)
+        except ValueError as error:
+            self.fail(f"generate_wordcloud raised ValueError unexpectedly: {error}")
+
+if __name__ == '__main__':
+    unittest.main()  # 执行测试