import unittest
from unittest.mock import patch, mock_open
import os
from 数据分析 import read_danmu, filter_ai_related_danmu, count_danmu, save, print_top_n_danmu

class TestAIDanmuAnalysis(unittest.TestCase):
    """
    测试 数据分析 模块的单元测试类。
    """

    def setUp(self):
        self.folder_path = 'test_data/弹幕收集按序'
        self.output_file = 'test_output.xlsx'

    def test_read_danmu(self):
        """
        测试读取弹幕文件。
        """
        test_data = "弹幕1\n弹幕2\n"
        with patch('builtins.open', new_callable=mock_open, read_data=test_data) as mock_file:
            danmu_list = read_danmu(self.folder_path)
            self.assertEqual(danmu_list, [('弹幕1', 'file1'), ('弹幕2', 'file1')])
            mock_file.assert_called_with(os.path.join(self.folder_path, 'file1.txt'), 'r', encoding='utf-8')

    def test_filter_ai_related_danmu(self):
        """
        测试筛选与 AI 相关的弹幕。
        """
        danmu_list = [('弹幕1', 'file1'), ('AI技术', 'file2')]
        filtered_danmu = filter_ai_related_danmu(danmu_list, ["AI"])
        self.assertEqual(filtered_danmu, [('AI技术', 'file2')])

    def test_count_danmu(self):
        """
        测试统计弹幕出现次数。
        """
        danmu_list = [('弹幕1', 'file1'), ('弹幕1', 'file2'), ('弹幕2', 'file1')]
        danmu_counter = count_danmu(danmu_list)
        self.assertEqual(danmu_counter, {'弹幕1': 2, '弹幕2': 1})

    def test_save(self):
        """
        测试将弹幕数据保存到 Excel 文件。
        """
        danmu_list = [('弹幕1', 'file1'), ('弹幕1', 'file2'), ('弹幕2', 'file1')]
        with patch('pandas.DataFrame.to_excel') as mock_to_excel:
            save(danmu_list, self.output_file)
            mock_to_excel.assert_called_with(self.output_file, index=False)

    def test_print_top_n_danmu(self):
        """
        测试打印数量排名前N的弹幕。
        """
        danmu_list = [('弹幕1', 'file1'), ('弹幕1', 'file2'), ('弹幕2', 'file1'), ('弹幕3', 'file1')]
        with patch('builtins.print') as mock_print:
            print_top_n_danmu(danmu_list, top_n=2)
            mock_print.assert_any_call('数量排名前 2 的弹幕：')
            mock_print.assert_any_call('弹幕内容: 弹幕1, 出现次数: 2')

if __name__ == '__main__':
    unittest.main()  # 执行测试