diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..6f1498d --- /dev/null +++ b/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 = ''' +
+ + ''' + 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 = '