| 1 | """LiveDownloader 测试。 |
| 2 | |
| 3 | 用 asyncio 伪造一个 aiohttp-like session，校验流选择、跳过非直播状态、成功路径。 |
| 4 | """ |
| 5 | |
| 6 | import asyncio |
| 7 | from pathlib import Path |
| 8 | |
| 9 | import pytest |
| 10 | |
| 11 | from auth import CookieManager |
| 12 | from config import ConfigLoader |
| 13 | from control import QueueManager, RateLimiter, RetryHandler |
| 14 | from core.api_client import DouyinAPIClient |
| 15 | from core.live_downloader import LiveDownloader |
| 16 | from storage import FileManager |
| 17 | |
| 18 | |
| 19 | class _FakeStreamResponse: |
| 20 | """模拟 aiohttp 流式响应。""" |
| 21 | |
| 22 | def __init__(self, chunks, status: int = 200): |
| 23 | self._chunks = list(chunks) |
| 24 | self.status = status |
| 25 | self.content = self |
| 26 | self._consumed = False |
| 27 | |
| 28 | async def __aenter__(self): |
| 29 | return self |
| 30 | |
| 31 | async def __aexit__(self, exc_type, exc, tb): |
| 32 | return False |
| 33 | |
| 34 | def iter_chunked(self, size): |
| 35 | async def gen(): |
| 36 | for c in self._chunks: |
| 37 | yield c |
| 38 | |
| 39 | return gen() |
| 40 | |
| 41 | |
| 42 | class _FakeSession: |
| 43 | def __init__(self, chunks, status: int = 200): |
| 44 | self._chunks = chunks |
| 45 | self._status = status |
| 46 | |
| 47 | def get(self, url, headers=None, timeout=None): |
| 48 | return _FakeStreamResponse(self._chunks, self._status) |
| 49 | |
| 50 | |
| 51 | def _build_downloader(tmp_path): |
| 52 | config = ConfigLoader() |
| 53 | config.update(path=str(tmp_path)) |
| 54 | |
| 55 | file_manager = FileManager(str(tmp_path)) |
| 56 | cookie_manager = CookieManager(str(tmp_path / ".cookies.json")) |
| 57 | api_client = DouyinAPIClient({}) |
| 58 | |
| 59 | return LiveDownloader( |
| 60 | config, |
| 61 | api_client, |
| 62 | file_manager, |
| 63 | cookie_manager, |
| 64 | database=None, |
| 65 | rate_limiter=RateLimiter(max_per_second=5), |
| 66 | retry_handler=RetryHandler(max_retries=1), |
| 67 | queue_manager=QueueManager(max_workers=1), |
| 68 | ), api_client |
| 69 | |
| 70 | |
| 71 | def test_select_best_stream_url_prefers_flv_origin(): |
| 72 | room = { |
| 73 | "stream_url": { |
| 74 | "flv_pull_url": { |
| 75 | "SD": "https://cdn/sd.flv", |
| 76 | "FULL_HD1": "https://cdn/fhd.flv", |
| 77 | "HD1": "https://cdn/hd.flv", |
| 78 | }, |
| 79 | "hls_pull_url_map": { |
| 80 | "HD1": "https://cdn/hd.m3u8", |
| 81 | }, |
| 82 | } |
| 83 | } |
| 84 | url, quality = LiveDownloader._select_best_stream_url(room) |
| 85 | assert url == "https://cdn/fhd.flv" |
| 86 | assert quality == "FULL_HD1" |
| 87 | |
| 88 | |
| 89 | def test_select_best_falls_back_to_hls(): |
| 90 | room = { |
| 91 | "stream_url": { |
| 92 | "hls_pull_url_map": {"HD1": "https://cdn/hd.m3u8", "SD": "https://cdn/sd.m3u8"}, |
| 93 | } |
| 94 | } |
| 95 | url, quality = LiveDownloader._select_best_stream_url(room) |
| 96 | assert url == "https://cdn/hd.m3u8" |
| 97 | assert quality == "HD1" |
| 98 | |
| 99 | |
| 100 | def test_select_best_returns_none_if_no_stream(): |
| 101 | assert LiveDownloader._select_best_stream_url({}) == (None, "") |
| 102 | assert LiveDownloader._select_best_stream_url({"stream_url": {}}) == (None, "") |
| 103 | |
| 104 | |
| 105 | @pytest.mark.asyncio |
| 106 | async def test_live_downloader_skips_when_not_live(tmp_path): |
| 107 | downloader, api_client = _build_downloader(tmp_path) |
| 108 | |
| 109 | async def fake_get_live_room_info(room_id, *, sec_user_id=""): |
| 110 | return {"room": {"status": 4, "stream_url": {}}, "user": {}} |
| 111 | |
| 112 | api_client.get_live_room_info = fake_get_live_room_info |
| 113 | |
| 114 | result = await downloader.download({"room_id": "42"}) |
| 115 | assert result.total == 1 |
| 116 | assert result.skipped == 1 |
| 117 | assert result.success == 0 |
| 118 | await api_client.close() |
| 119 | |
| 120 | |
| 121 | @pytest.mark.asyncio |
| 122 | async def test_live_downloader_records_stream(tmp_path, monkeypatch): |
| 123 | downloader, api_client = _build_downloader(tmp_path) |
| 124 | |
| 125 | async def fake_get_live_room_info(room_id, *, sec_user_id=""): |
| 126 | return { |
| 127 | "room": { |
| 128 | "status": 2, |
| 129 | "title": "测试标题", |
| 130 | "stream_url": { |
| 131 | "flv_pull_url": {"ORIGIN": "https://cdn/live.flv"}, |
| 132 | }, |
| 133 | }, |
| 134 | "user": {"nickname": "主播甲"}, |
| 135 | } |
| 136 | |
| 137 | api_client.get_live_room_info = fake_get_live_room_info |
| 138 | |
| 139 | async def fake_get_session(): |
| 140 | return _FakeSession([b"abc", b"def", b"ghi"]) |
| 141 | |
| 142 | api_client.get_session = fake_get_session |
| 143 | |
| 144 | result = await downloader.download({"room_id": "42"}) |
| 145 | assert result.success == 1 |
| 146 | # 找到已落盘的 flv 文件 |
| 147 | flvs = list(tmp_path.rglob("*.flv")) |
| 148 | assert len(flvs) == 1 |
| 149 | assert flvs[0].read_bytes() == b"abcdefghi" |
| 150 | |
| 151 | await api_client.close() |
| 152 | |
| 153 | |
| 154 | @pytest.mark.asyncio |
| 155 | async def test_live_downloader_reports_failure_on_missing_stream(tmp_path): |
| 156 | downloader, api_client = _build_downloader(tmp_path) |
| 157 | |
| 158 | async def fake_info(room_id, *, sec_user_id=""): |
| 159 | return {"room": {"status": 2, "stream_url": {}}, "user": {}} |
| 160 | |
| 161 | api_client.get_live_room_info = fake_info |
| 162 | result = await downloader.download({"room_id": "42"}) |
| 163 | assert result.failed == 1 |
| 164 | await api_client.close() |
| 165 | |
| 166 | |
| 167 | @pytest.mark.asyncio |
| 168 | async def test_live_downloader_fails_when_room_missing(tmp_path): |
| 169 | downloader, api_client = _build_downloader(tmp_path) |
| 170 | |
| 171 | async def fake_info(room_id, *, sec_user_id=""): |
| 172 | return None |
| 173 | |
| 174 | api_client.get_live_room_info = fake_info |
| 175 | result = await downloader.download({"room_id": "42"}) |
| 176 | assert result.failed == 1 |
| 177 | await api_client.close() |
| 178 | |
| 179 | |
| 180 | class _IdleTimeoutSession: |
| 181 | """模拟写入一些数据后抛 asyncio.TimeoutError 的流。""" |
| 182 | |
| 183 | def __init__(self, chunks_before_timeout): |
| 184 | self._chunks = chunks_before_timeout |
| 185 | |
| 186 | def get(self, url, headers=None, timeout=None): |
| 187 | chunks = self._chunks |
| 188 | |
| 189 | class _Resp: |
| 190 | status = 200 |
| 191 | |
| 192 | async def __aenter__(self_inner): |
| 193 | return self_inner |
| 194 | |
| 195 | async def __aexit__(self_inner, exc_type, exc, tb): |
| 196 | return False |
| 197 | |
| 198 | @property |
| 199 | def content(self_inner): |
| 200 | return self_inner |
| 201 | |
| 202 | def iter_chunked(self_inner, size): |
| 203 | async def gen(): |
| 204 | for c in chunks: |
| 205 | yield c |
| 206 | raise asyncio.TimeoutError("sock_read idle") |
| 207 | |
| 208 | return gen() |
| 209 | |
| 210 | return _Resp() |
| 211 | |
| 212 | |
| 213 | @pytest.mark.asyncio |
| 214 | async def test_live_downloader_preserves_partial_on_idle_timeout(tmp_path): |
| 215 | """idle_timeout（主播停止推流 / 网络卡住）不应丢弃已录制的字节。""" |
| 216 | downloader, api_client = _build_downloader(tmp_path) |
| 217 | |
| 218 | async def fake_info(room_id, *, sec_user_id=""): |
| 219 | return { |
| 220 | "room": { |
| 221 | "status": 2, |
| 222 | "title": "测试", |
| 223 | "stream_url": {"flv_pull_url": {"ORIGIN": "https://cdn/live.flv"}}, |
| 224 | }, |
| 225 | "user": {"nickname": "主播"}, |
| 226 | } |
| 227 | |
| 228 | api_client.get_live_room_info = fake_info |
| 229 | |
| 230 | async def fake_get_session(): |
| 231 | return _IdleTimeoutSession([b"partial1", b"partial2"]) |
| 232 | |
| 233 | api_client.get_session = fake_get_session |
| 234 | |
| 235 | result = await downloader.download({"room_id": "42"}) |
| 236 | assert result.success == 1 # 保留部分数据视为成功 |
| 237 | flvs = list(tmp_path.rglob("*.flv")) |
| 238 | assert len(flvs) == 1 |
| 239 | assert flvs[0].read_bytes() == b"partial1partial2" |
| 240 | # 没有遗留 .tmp 文件 |
| 241 | tmps = list(tmp_path.rglob("*.tmp")) |
| 242 | assert tmps == [] |
| 243 | |
| 244 | await api_client.close() |
| 245 | |
| 246 | |
| 247 | def test_download_headers_referer_for_live(tmp_path): |
| 248 | """_record_stream 内部应把 Referer 改为 live.douyin.com（通过构造样本流间接验证）。""" |
| 249 | # 直接读源码断言是更轻量的方式——避免走完整集成路径。 |
| 250 | |
| 251 | source = Path("core/live_downloader.py").read_text(encoding="utf-8") |
| 252 | assert 'headers["Referer"] = "https://live.douyin.com/"' in source |
| 253 | |
| 254 | |
| 255 | def test_download_headers_origin_for_live(tmp_path): |
| 256 | """_record_stream 应同时改写 Origin 头（CDN 可能同时校验 Referer 与 Origin）。""" |
| 257 | |
| 258 | source = Path("core/live_downloader.py").read_text(encoding="utf-8") |
| 259 | assert 'headers["Origin"] = "https://live.douyin.com"' in source |
| 260 | |
| 261 | |
| 262 | class _HeaderCapturingSession: |
| 263 | """捕获 get() 调用时传入的 headers，供断言。""" |
| 264 | |
| 265 | def __init__(self, chunks): |
| 266 | self._chunks = chunks |
| 267 | self.captured_headers = None |
| 268 | |
| 269 | def get(self, url, headers=None, timeout=None): |
| 270 | self.captured_headers = dict(headers or {}) |
| 271 | chunks = self._chunks |
| 272 | |
| 273 | class _Resp: |
| 274 | status = 200 |
| 275 | |
| 276 | async def __aenter__(self_inner): |
| 277 | return self_inner |
| 278 | |
| 279 | async def __aexit__(self_inner, *_args): |
| 280 | return False |
| 281 | |
| 282 | @property |
| 283 | def content(self_inner): |
| 284 | return self_inner |
| 285 | |
| 286 | def iter_chunked(self_inner, size): |
| 287 | async def gen(): |
| 288 | for c in chunks: |
| 289 | yield c |
| 290 | |
| 291 | return gen() |
| 292 | |
| 293 | return _Resp() |
| 294 | |
| 295 | |
| 296 | @pytest.mark.asyncio |
| 297 | async def test_live_recording_sends_live_origin_and_referer(tmp_path): |
| 298 | """端到端验证：_record_stream 实际发出的请求头含正确的 live.douyin.com。""" |
| 299 | downloader, api_client = _build_downloader(tmp_path) |
| 300 | |
| 301 | async def fake_info(room_id, *, sec_user_id=""): |
| 302 | return { |
| 303 | "room": { |
| 304 | "status": 2, |
| 305 | "title": "t", |
| 306 | "stream_url": {"flv_pull_url": {"ORIGIN": "https://cdn/live.flv"}}, |
| 307 | }, |
| 308 | "user": {"nickname": "主播"}, |
| 309 | } |
| 310 | |
| 311 | api_client.get_live_room_info = fake_info |
| 312 | capturing = _HeaderCapturingSession([b"data"]) |
| 313 | |
| 314 | async def fake_get_session(): |
| 315 | return capturing |
| 316 | |
| 317 | api_client.get_session = fake_get_session |
| 318 | |
| 319 | await downloader.download({"room_id": "42"}) |
| 320 | assert capturing.captured_headers is not None |
| 321 | assert capturing.captured_headers.get("Referer") == "https://live.douyin.com/" |
| 322 | assert capturing.captured_headers.get("Origin") == "https://live.douyin.com" |
| 323 | |
| 324 | await api_client.close() |
| 325 |