homework/test_script_2_1.py

# import unittest
# from vector import Vector
import unittest
from unittest.mock import patch, MagicMock
import requests

# 导入待测试的函数，从 script_2_1 文件中导入
from script_2_1 import get_search_results, get_cid, get_danmaku, main

class TestBilibiliDanmaku(unittest.TestCase):

    @patch('requests.get')
    def test_get_search_results(self, mock_get):
        # 模拟API返回的数据
        mock_response = MagicMock()
        mock_response.json.return_value = {
            'data': {
                'result': [
                    {'bvid': 'BV1xx411c7a6', 'title': 'Test Video 1'},
                    {'bvid': 'BV2xx411c7a7', 'title': 'Test Video 2'}
                ]
            }
        }
        mock_response.status_code = 200
        mock_get.return_value = mock_response
        
        keyword = "2024巴黎奥运会"
        page = 1
        
        # 调用被测试的函数
        result = get_search_results(keyword, page)
        
        # 检查API返回是否正确
        self.assertEqual(len(result), 2)
        self.assertEqual(result[0]['bvid'], 'BV1xx411c7a6')
        self.assertEqual(result[1]['title'], 'Test Video 2')

    @patch('requests.get')
    def test_get_cid(self, mock_get):
        # 模拟返回的CID信息
        mock_response = MagicMock()
        mock_response.json.return_value = {
            'data': [
                {'cid': 123456}
            ]
        }
        mock_response.status_code = 200
        mock_get.return_value = mock_response
        
        bvid = "BV1xx411c7a6"
        cid = get_cid(bvid)
        
        # 检查CID是否正确
        self.assertEqual(cid, 123456)

    @patch('requests.get')
    def test_get_danmaku(self, mock_get):
        # 模拟弹幕XML
        mock_response = MagicMock()
        mock_response.content = '''<i><d p="...">弹幕内容1</d><d p="...">弹幕内容2</d></i>'''.encode('utf-8')
        mock_response.status_code = 200
        mock_get.return_value = mock_response
        
        cid = 123456
        bvid = "BV1xx411c7a6"
        danmakus = get_danmaku(cid, bvid)
        
        # 检查弹幕内容
        self.assertEqual(len(danmakus), 2)
        self.assertIn('弹幕内容1', danmakus)
        self.assertIn('弹幕内容2', danmakus)

    @patch('requests.get')
    @patch('script_2_1.get_cid')
    @patch('script_2_1.get_danmaku')
    def test_main(self, mock_get_danmaku, mock_get_cid, mock_get):
        # 模拟搜索视频接口的返回数据
        mock_response = MagicMock()
        mock_response.json.return_value = {
            'data': {
                'result': [
                    {'bvid': 'BV1xx411c7a6', 'title': 'Test Video 1'},
                    {'bvid': 'BV2xx411c7a7', 'title': 'Test Video 2'}
                ]
            }
        }
        mock_response.status_code = 200
        mock_get.return_value = mock_response

        # 模拟 get_cid 的返回值
        mock_get_cid.return_value = 123456

        # 模拟 get_danmaku 的返回值
        mock_get_danmaku.return_value = ['弹幕内容1', '弹幕内容2']

        # 运行主函数
        with patch('builtins.open', unittest.mock.mock_open()) as mock_file:
            main()

        # 确保打开文件被调用
        mock_file.assert_called_once_with('danmakus_2024_olympics.txt', 'w', encoding='utf-8')
        handle = mock_file()
        handle.write.assert_any_call('弹幕内容1\n')
        handle.write.assert_any_call('弹幕内容2\n')


if __name__ == '__main__':
    unittest.main()