diff --git a/app/modules/themoviedb/tmdbv3api/tmdb.py b/app/modules/themoviedb/tmdbv3api/tmdb.py index aa4fc0777..e03acfb23 100644 --- a/app/modules/themoviedb/tmdbv3api/tmdb.py +++ b/app/modules/themoviedb/tmdbv3api/tmdb.py @@ -56,6 +56,8 @@ def _is_empty_result_snapshot(snapshot) -> bool: class TMDb(object): + """封装 TMDB 请求、响应缓存及有界的连接和限流重试。""" + _RESPONSE_SNAPSHOT_MARKER = "__mp_tmdb_response_snapshot__" def __init__(self, session=None, language=None): @@ -365,26 +367,21 @@ class TMDb(object): ) def _handle_headers(self, headers): + """只记录当前响应的限流状态，缺失字段不能沿用上次耗尽的配额。""" normalized_headers = { str(key).lower(): value for key, value in dict(headers or {}).items() } - if "x-ratelimit-remaining" in normalized_headers: - self._remaining = int(normalized_headers["x-ratelimit-remaining"]) - - if "x-ratelimit-reset" in normalized_headers: - self._reset = int(normalized_headers["x-ratelimit-reset"]) + self._remaining = int(normalized_headers.get("x-ratelimit-remaining", 40)) + reset = normalized_headers.get("x-ratelimit-reset") + self._reset = int(reset) if reset is not None else None def _handle_rate_limit(self): + """计算尚未到期的限流等待；缺少重置时间时直接报告限流。""" if self._remaining < 1: - current_time = int(time.time()) - sleep_time = self._reset - current_time - - if self.wait_on_rate_limit: - logger.warning("达到请求频率限制，休眠：%d 秒..." % sleep_time) - return abs(sleep_time) - else: + if not self.wait_on_rate_limit or self._reset is None: raise TMDbException("达到请求频率限制，请稍后再试！") + return max(0, self._reset - int(time.time())) return 0 def _process_json_response(self, json_data, is_async=False): @@ -419,23 +416,30 @@ class TMDb(object): def _request_obj(self, action, params="", call_cached=True, method="GET", data=None, json=None, key=None): + """同步解析响应；限流最多等待重试一次，避免递归耗尽请求调用栈。""" self._validate_api_key() url = self._build_url(action, params) - with fresh(not call_cached or method == "POST"): - req = self.request(method, url, data, json, - _ts=datetime.strftime(datetime.now(), '%Y%m%d')) + rate_limit_retried = False + while True: + with fresh(not call_cached or method == "POST"): + req = self.request(method, url, data, json, + _ts=datetime.strftime(datetime.now(), '%Y%m%d')) - if req is None: - return None + if req is None: + return None - self._handle_headers(self._get_response_headers(req)) + self._handle_headers(self._get_response_headers(req)) - rate_limit_result = self._handle_rate_limit() - if rate_limit_result: - logger.warning("达到请求频率限制，将在 %d 秒后重试..." % rate_limit_result) + rate_limit_result = self._handle_rate_limit() + if not rate_limit_result: + break + if rate_limit_retried: + raise TMDbException("达到请求频率限制，请稍后再试！") + logger.warning(f"达到请求频率限制，将在 {rate_limit_result} 秒后重试...") time.sleep(rate_limit_result) - return self._request_obj(action, params, False, method, data, json, key) + rate_limit_retried = True + call_cached = False json_data = self._get_response_json(req) self._validate_json_response(json_data) @@ -448,23 +452,30 @@ class TMDb(object): async def _async_request_obj(self, action, params="", call_cached=True, method="GET", data=None, json=None, key=None): + """异步解析响应，沿用同步入口的单次限流重试上限。""" self._validate_api_key() url = self._build_url(action, params) - async with async_fresh(not call_cached or method == "POST"): - req = await self.async_request(method, url, data, json, - _ts=datetime.strftime(datetime.now(), '%Y%m%d')) + rate_limit_retried = False + while True: + async with async_fresh(not call_cached or method == "POST"): + req = await self.async_request(method, url, data, json, + _ts=datetime.strftime(datetime.now(), '%Y%m%d')) - if req is None: - return None + if req is None: + return None - self._handle_headers(self._get_response_headers(req)) + self._handle_headers(self._get_response_headers(req)) - rate_limit_result = self._handle_rate_limit() - if rate_limit_result: - logger.warning("达到请求频率限制，将在 %d 秒后重试..." % rate_limit_result) + rate_limit_result = self._handle_rate_limit() + if not rate_limit_result: + break + if rate_limit_retried: + raise TMDbException("达到请求频率限制，请稍后再试！") + logger.warning(f"达到请求频率限制，将在 {rate_limit_result} 秒后重试...") await asyncio.sleep(rate_limit_result) - return await self._async_request_obj(action, params, False, method, data, json, key) + rate_limit_retried = True + call_cached = False json_data = self._get_response_json(req) self._validate_json_response(json_data) diff --git a/tests/fixtures/architecture/mypy-baseline.json b/tests/fixtures/architecture/mypy-baseline.json index 3ad5d25cc..5e47f8a8c 100644 --- a/tests/fixtures/architecture/mypy-baseline.json +++ b/tests/fixtures/architecture/mypy-baseline.json @@ -2248,9 +2248,8 @@ "assignment": 1, "attr-defined": 3, "no-redef": 2, - "no-untyped-call": 36, - "no-untyped-def": 44, - "operator": 1 + "no-untyped-call": 34, + "no-untyped-def": 44 }, "app/modules/thetvdb/__init__.py": { "empty-body": 1, diff --git a/tests/test_tmdb_rate_limit.py b/tests/test_tmdb_rate_limit.py new file mode 100644 index 000000000..0e3b0beea --- /dev/null +++ b/tests/test_tmdb_rate_limit.py @@ -0,0 +1,159 @@ +"""TMDB 限流恢复必须有界，且只使用当前响应的限流状态。""" + +import asyncio +from unittest.mock import AsyncMock, Mock + +import pytest +import requests + +from app.modules.themoviedb.tmdbv3api import tmdb as tmdb_module +from app.modules.themoviedb.tmdbv3api.exceptions import TMDbConnectionError, TMDbException +from app.modules.themoviedb.tmdbv3api.tmdb import TMDb + + +@pytest.fixture(params=[False, True], ids=['sync', 'async']) +def client_case(request, monkeypatch): + """同时驱动真实同步、异步缓存入口，只替换单次传输与等待。""" + client = TMDb() + client.api_key = 'test-key' + client.domain = 'tmdb.example.test' + transport = AsyncMock() if request.param else Mock() + sleep = AsyncMock() if request.param else Mock() + monkeypatch.setattr(tmdb_module.time, 'time', lambda: 1000) + if request.param: + monkeypatch.setattr(client, '_async_request_once', transport) + monkeypatch.setattr(tmdb_module.asyncio, 'sleep', sleep) + else: + monkeypatch.setattr(client, '_request_once', transport) + monkeypatch.setattr(tmdb_module.time, 'sleep', sleep) + + def invoke(method='GET'): + """在同一断言中调用同步或异步的完整响应解析入口。""" + if request.param: + return asyncio.run(client._async_request_obj('/movie/42', method=method)) + return client._request_obj('/movie/42', method=method) + + client.cache_clear() + asyncio.run(client.async_request.cache_clear()) + yield client, transport, sleep, invoke + client.cache_clear() + asyncio.run(client.async_request.cache_clear()) + client.close() + + +def _response(headers): + """返回合法成功响应，保证失败来自限流处理而非响应解析。""" + return Mock(headers=headers, json=Mock(return_value={'id': 42})) + + +@pytest.mark.parametrize('method', ['GET', 'POST']) +@pytest.mark.parametrize('recovered_headers', [{}, {'X-RateLimit-Remaining': '39'}]) +def test_rate_limit_recovery_uses_current_headers(client_case, method, recovered_headers): + """重试绕过旧缓存，且成功响应缺少限流头时不继承上次耗尽状态。""" + _, transport, sleep, invoke = client_case + transport.side_effect = [ + _response({'X-RateLimit-Remaining': '0', 'X-RateLimit-Reset': '1002'}), + _response(recovered_headers), + ] + + assert invoke(method) == {'id': 42} + assert transport.call_count == 2 + sleep.assert_called_once_with(2) + + +def test_persistent_rate_limit_stops_after_one_retry(client_case): + """持续限流至多重试一次，避免递归最终在 Cookie 等深层调用中爆栈。""" + _, transport, sleep, invoke = client_case + transport.return_value = _response( + {'X-RateLimit-Remaining': '0', 'X-RateLimit-Reset': '1002'} + ) + + with pytest.raises(TMDbException, match='达到请求频率限制') as error: + invoke() + + assert not isinstance(error.value, TMDbConnectionError) + assert transport.call_count == 2 + sleep.assert_called_once_with(2) + + +@pytest.mark.parametrize('reset', ['998', '1000']) +def test_expired_rate_limit_does_not_sleep(client_case, reset): + """已到期的重置时间不能取绝对值后继续等待和重试。""" + _, transport, sleep, invoke = client_case + transport.return_value = _response( + {'X-RateLimit-Remaining': '0', 'X-RateLimit-Reset': reset} + ) + + assert invoke() == {'id': 42} + assert transport.call_count == 1 + sleep.assert_not_called() + + +@pytest.mark.parametrize('wait_enabled', [False, True]) +def test_incomplete_rate_limit_has_no_guessed_wait(client_case, wait_enabled): + """没有重置时间时报告限流，避免沿用旧时间或发生 None 减法异常。""" + client, transport, sleep, invoke = client_case + client.wait_on_rate_limit = wait_enabled + client._reset = 1002 + transport.return_value = _response({'X-RateLimit-Remaining': '0'}) + + with pytest.raises(TMDbException, match='达到请求频率限制'): + invoke() + + assert transport.call_count == 1 + sleep.assert_not_called() + + +def test_disabled_wait_reports_rate_limit_immediately(client_case): + """调用方关闭限流等待时，不执行重试或休眠。""" + client, transport, sleep, invoke = client_case + client.wait_on_rate_limit = False + transport.return_value = _response( + {'X-RateLimit-Remaining': '0', 'X-RateLimit-Reset': '1002'} + ) + + with pytest.raises(TMDbException, match='达到请求频率限制'): + invoke() + + assert transport.call_count == 1 + sleep.assert_not_called() + + +def test_real_cookie_session_recovers_after_rate_limit(monkeypatch): + """保留真实 Session/Cookie 准备，验证缺失限流头的代理响应不会引发递归。""" + sent = [] + + def send(_adapter, request, **kwargs): + """仅替换网络发送，依次返回耗尽配额与无配额头的成功响应。""" + sent.append(request) + assert request.headers['Cookie'] == 'probe=present' + assert kwargs['proxies']['https'] == 'http://proxy.example.test:8080' + response = requests.Response() + response.request = request + response.url = request.url + response.status_code = 200 + response._content = b'{"id": 42}' + if len(sent) == 1: + response.headers.update({ + 'X-RateLimit-Remaining': '0', + 'X-RateLimit-Reset': '1002', + }) + return response + + monkeypatch.setattr(requests.adapters.HTTPAdapter, 'send', send) + monkeypatch.setattr(tmdb_module.time, 'time', lambda: 1000) + sleep = Mock() + monkeypatch.setattr(tmdb_module.time, 'sleep', sleep) + with requests.Session() as session: + session.trust_env = False + session.cookies.set('probe', 'present', domain='tmdb.example.test', path='/') + client = TMDb(session=session) + client.api_key = 'test-key' + client.domain = 'tmdb.example.test' + client._req._proxies = {'https': 'http://proxy.example.test:8080'} + try: + assert client._request_obj('/movie/42', call_cached=False) == {'id': 42} + assert len(sent) == 2 + sleep.assert_called_once_with(2) + finally: + client.cache_clear()