""" 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() # 执行测试