"""
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()  # 执行测试