对script_2_1进行单元测试

test_script_2_1.py

@@ -0,0 +1,108 @@
+# 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()