102201539/tests/test_main.py

"""
test_main.py - 单元测试主模块的功能。
"""
import unittest
from unittest.mock import patch, MagicMock
import pandas as pd
from main import (
    search_videos,
    get_cid,
    get_danmaku,
    parse_danmaku,
    save_danmakus_to_excel
)

class TestMain(unittest.TestCase):
    """
    测试主模块的单元测试类。
    """

    @patch('main.requests.get')
    def test_search_videos_success(self, mock_get):
        """
        测试根据关键词搜索视频并成功返回视频 ID 列表。
        """
        # 模拟返回的 HTML 内容
        mock_html = '''
        <div class="bili-video-card">
            <a href="//www.bilibili.com/video/123456/"></a>
        </div>
        <div class="bili-video-card">
            <a href="//www.bilibili.com/video/789012/"></a>
        </div>       
        '''
        mock_get.return_value = MagicMock(status_code=200, text=mock_html)

        # 调用搜索函数
        video_ids = search_videos("2024巴黎奥运会", max_results=2)
        # 期望返回的结果
        self.assertEqual(video_ids, ['123456', '789012'])

    @patch('main.requests.get')
    def test_search_videos_failure(self, mock_get):
        """
        测试根据关键词搜索视频时请求失败。
        """
        mock_get.return_value = MagicMock(status_code=404)

        # 调用搜索函数
        video_ids = search_videos("2024巴黎奥运会", max_results=2)
        # 期望返回空列表
        self.assertEqual(video_ids, [])

    @patch('main.requests.get')
    def test_get_cid_success(self, mock_get):
        """
        测试根据视频 ID 获取弹幕 CID 成功。
        """
        mock_html = '{"cid":123456}'
        mock_get.return_value = MagicMock(status_code=200, text=mock_html)

        cid = get_cid("123456")
        # 期望返回的 CID
        self.assertEqual(cid, '123456')

    @patch('main.requests.get')
    def test_get_cid_failure(self, mock_get):
        """
        测试根据视频 ID 获取弹幕 CID 失败。
        """
        mock_get.return_value = MagicMock(status_code=404)

        cid = get_cid("123456")
        # 期望返回 None
        self.assertIsNone(cid)

    @patch('main.requests.get')
    def test_get_danmaku_success(self, mock_get):
        """
        测试根据视频 ID 和 CID 获取弹幕数据成功。
        """
        mock_danmaku_xml = '<i><d p="0,1,25,16777215,0,0,0,0">弹幕内容</d></i>'
        mock_get.return_value = MagicMock(status_code=200, content=mock_danmaku_xml.encode('utf-8'))

        danmaku = get_danmaku("123456", "123456")
        # 期望返回的弹幕 XML 内容
        self.assertEqual(danmaku, mock_danmaku_xml.encode('utf-8'))

    @patch('main.requests.get')
    def test_get_danmaku_failure(self, mock_get):
        """
        测试根据视频 ID 和 CID 获取弹幕数据失败。
        """
        mock_get.return_value = MagicMock(status_code=404)

        danmaku = get_danmaku("123456", "123456")
        # 期望返回 None
        self.assertIsNone(danmaku)

    def test_parse_danmaku(self):
        """
        测试解析弹幕 XML 内容，提取弹幕文本。
        """
        xml_content = '<i><d p="0,1,25,16777215,0,0,0,0">弹幕内容1</d><d \
            p="0,1,25,16777215,0,0,0,0">弹幕内容2</d></i>'
        result = parse_danmaku(xml_content)
        # 期望返回的弹幕列表
        self.assertEqual(result, ['弹幕内容1', '弹幕内容2'])

    def test_save_danmakus_to_excel(self):
        """
        测试将弹幕数据保存到 Excel 文件中。
        """
        all_danmakus = ['弹幕内容1', '弹幕内容2']
        filename = "test_danmakus.xlsx"

        # 调用保存函数
        save_danmakus_to_excel(all_danmakus, filename)

        # 读取保存的文件并验证内容
        data_frame = pd.read_excel(filename)
        self.assertEqual(list(data_frame['弹幕内容']), all_danmakus)

    def test_save_danmakus_to_excel_empty(self):
        """
        测试保存空的弹幕列表到 Excel 文件的功能。
        """
        all_danmakus = []
        filename = "empty_danmakus.xlsx"

        # 调用保存函数
        save_danmakus_to_excel(all_danmakus, filename)

        # 验证保存的文件是否为空
        data_frame = pd.read_excel(filename)
        self.assertEqual(len(data_frame), 0)

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