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 = '弹幕内容' + 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 = '弹幕内容1弹幕内容2' + 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() # 执行测试