ADD file via upload

tests/test_main.py

@@ -0,0 +1,138 @@
+"""
+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()  # 执行测试