diff --git a/src/tiktok_research_api_helper/api_client.py b/src/tiktok_research_api_helper/api_client.py index 6cd52cd..63670b6 100644 --- a/src/tiktok_research_api_helper/api_client.py +++ b/src/tiktok_research_api_helper/api_client.py @@ -63,7 +63,7 @@ class TiktokCredentials: @attrs.define -class TikTokResponse: +class TikTokVideoResponse: data: Mapping[str, Any] videos: Sequence[Any] error: Mapping[str, Any] @@ -400,15 +400,15 @@ def _fetch_retryer(self, max_api_rate_limit_retries=None): reraise=True, ) - def fetch(self, request: TiktokRequest, max_api_rate_limit_retries=None) -> TikTokResponse: + def fetch(self, request: TiktokRequest, max_api_rate_limit_retries=None) -> TikTokVideoResponse: return self._fetch_retryer(max_api_rate_limit_retries=max_api_rate_limit_retries)( self._fetch, request ) - def _fetch(self, request: TiktokRequest) -> TikTokResponse: + def _fetch(self, request: TiktokRequest) -> TikTokVideoResponse: api_response = self._post(request) self._num_api_requests_sent += 1 - return self._parse_response(api_response) + return self._parse_video_response(api_response) @tenacity.retry( stop=tenacity.stop_after_attempt(10), @@ -458,30 +458,33 @@ def _post(self, request: TiktokRequest) -> rq.Response | None: return None @staticmethod - def _parse_response(response: rq.Response | None) -> TikTokResponse: - if response is None: - raise ValueError("Response is None") + def _parse_video_response(response: rq.Response | None) -> TikTokVideoResponse: + response_json = _extract_response_json(response) + error_data = response_json.get("error") + response_data_section = response_json.get("data", {}) + videos = response_data_section.get("videos", []) - try: - response_json = response.json() - response_data_section = response_json.get("data", {}) - error_data = response_json.get("error") - except rq.exceptions.JSONDecodeError: - logging.info( - "Error parsing JSON response:\n%s\n%s\n%s", - response.status_code, - "\n".join([f"{k}: {v}" for k, v in response.headers.items()]), - response.text, - ) - raise + return TikTokVideoResponse(data=response_data_section, videos=videos, error=error_data) - videos = response_data_section.get("videos", []) - return TikTokResponse(data=response_data_section, videos=videos, error=error_data) +def _extract_response_json(response: rq.Response | None) -> Mapping[str, Any]: + if response is None: + raise ValueError("Response is None") + + try: + return response.json() + except rq.exceptions.JSONDecodeError: + logging.info( + "Error parsing JSON response:\n%s\n%s\n%s", + response.status_code, + "\n".join([f"{k}: {v}" for k, v in response.headers.items()]), + response.text, + ) + raise def update_crawl_from_api_response( - crawl: Crawl, api_response: TikTokResponse, num_videos_requested: int = 100 + crawl: Crawl, api_response: TikTokVideoResponse, num_videos_requested: int = 100 ): crawl.cursor = api_response.data["cursor"] crawl.has_more = api_response.data["has_more"] diff --git a/tests/test_api_client.py b/tests/test_api_client.py index 925d977..4af96a2 100644 --- a/tests/test_api_client.py +++ b/tests/test_api_client.py @@ -233,17 +233,17 @@ def mock_tiktok_responses(testdata_api_response_json): video["id"] += 1 return [ - api_client.TikTokResponse( + api_client.TikTokVideoResponse( data=first_page["data"], videos=first_page["data"]["videos"], error=first_page["error"], ), - api_client.TikTokResponse( + api_client.TikTokVideoResponse( data=second_page["data"], videos=second_page["data"]["videos"], error=second_page["error"], ), - api_client.TikTokResponse( + api_client.TikTokVideoResponse( data=last_page["data"], videos=last_page["data"]["videos"], error=last_page["error"],