From fae3cd9801ef1fbe865a51d0fb5781d8b0f62690 Mon Sep 17 00:00:00 2001 From: wen Date: Wed, 4 Dec 2024 13:41:45 +0800 Subject: [PATCH 1/2] fix: fix crawl pixiv --- src/favorites_crawler/commands/crawl.py | 5 ++++- src/favorites_crawler/middlewares.py | 8 +------- src/favorites_crawler/utils/auth.py | 6 +++--- 3 files changed, 8 insertions(+), 11 deletions(-) diff --git a/src/favorites_crawler/commands/crawl.py b/src/favorites_crawler/commands/crawl.py index 3170f50..bf16250 100644 --- a/src/favorites_crawler/commands/crawl.py +++ b/src/favorites_crawler/commands/crawl.py @@ -10,6 +10,7 @@ from favorites_crawler.utils.config import load_config, overwrite_spider_settings from favorites_crawler.constants.path import DEFAULT_FAVORS_HOME +from favorites_crawler.utils.auth import refresh_pixiv_token app = typer.Typer(help='Crawl your favorites from websites.', no_args_is_help=True) @@ -27,7 +28,9 @@ def crawl_yandere(): @app.command('pixiv') def crawl_pixiv(): """Crawl your favorite illustrations from pixiv.""" - crawl('pixiv') + favors_home = os.getenv('FAVORS_HOME', DEFAULT_FAVORS_HOME) + access_token = refresh_pixiv_token(favors_home) + crawl('pixiv', access_token=access_token) @app.command('nhentai') diff --git a/src/favorites_crawler/middlewares.py b/src/favorites_crawler/middlewares.py index 9676435..d150c6f 100644 --- a/src/favorites_crawler/middlewares.py +++ b/src/favorites_crawler/middlewares.py @@ -3,13 +3,7 @@ # See documentation in: # https://docs.scrapy.org/en/latest/topics/spider-middleware.html -from favorites_crawler.utils.auth import refresh_pixiv - class PixivAuthorizationMiddleware: - def __init__(self): - self.access_token = refresh_pixiv() - def process_request(self, request, spider): - if self.access_token: - request.headers.setdefault(b'Authorization', f'Bearer {self.access_token}') + request.headers.setdefault(b'Authorization', f'Bearer {spider.access_token}') diff --git a/src/favorites_crawler/utils/auth.py b/src/favorites_crawler/utils/auth.py index 85762f6..c598e6b 100644 --- a/src/favorites_crawler/utils/auth.py +++ b/src/favorites_crawler/utils/auth.py @@ -23,8 +23,8 @@ def _GetPixivToken__wait_for_redirect(self) -> None: raise ValueError(msg) from err -def refresh_pixiv(): - config = load_config() +def refresh_pixiv_token(home: str): + config = load_config(home) pixiv_config = config.get('pixiv', {}) refresh_token = pixiv_config.get('REFRESH_TOKEN') if not refresh_token: @@ -33,7 +33,7 @@ def refresh_pixiv(): login_info = token_getter.refresh(refresh_token) access_token = login_info['access_token'] pixiv_config['ACCESS_TOKEN'] = access_token - dump_config(config) + dump_config(config, home) return access_token From 82006d1b2b8387fa9e4f83283a971ee2578327cf Mon Sep 17 00:00:00 2001 From: wen Date: Wed, 4 Dec 2024 14:03:23 +0800 Subject: [PATCH 2/2] test: test refresh_pixiv_token --- src/favorites_crawler/utils/auth.py | 5 +++- src/favorites_crawler/utils/config.py | 9 ++++--- tests/test_utils/test_auth.py | 38 ++++++++++++++++++++++++++- 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/src/favorites_crawler/utils/auth.py b/src/favorites_crawler/utils/auth.py index c598e6b..cb32608 100644 --- a/src/favorites_crawler/utils/auth.py +++ b/src/favorites_crawler/utils/auth.py @@ -1,5 +1,8 @@ +from __future__ import annotations + import json import re +from pathlib import Path from urllib.parse import unquote from gppt import GetPixivToken @@ -23,7 +26,7 @@ def _GetPixivToken__wait_for_redirect(self) -> None: raise ValueError(msg) from err -def refresh_pixiv_token(home: str): +def refresh_pixiv_token(home: str | Path): config = load_config(home) pixiv_config = config.get('pixiv', {}) refresh_token = pixiv_config.get('REFRESH_TOKEN') diff --git a/src/favorites_crawler/utils/config.py b/src/favorites_crawler/utils/config.py index d4ab627..55f079e 100644 --- a/src/favorites_crawler/utils/config.py +++ b/src/favorites_crawler/utils/config.py @@ -1,4 +1,7 @@ +from __future__ import annotations + import os +from pathlib import Path from copy import deepcopy import yaml @@ -37,7 +40,7 @@ } -def load_config(home: str) -> dict: +def load_config(home: str | Path) -> dict: """Load config from user home""" home = os.path.expanduser(home) create_favors_home(home) @@ -49,7 +52,7 @@ def load_config(home: str) -> dict: return yaml.safe_load(f) -def dump_config(data: dict, home: str): +def dump_config(data: dict, home: str | Path): """Dump config data to user home""" home = os.path.expanduser(home) create_favors_home(home) @@ -58,7 +61,7 @@ def dump_config(data: dict, home: str): yaml.safe_dump(data, f, allow_unicode=True) -def create_favors_home(path: str): +def create_favors_home(path: str | Path): """Create favors home if not exists""" if not os.path.exists(path): os.makedirs(path, exist_ok=True) diff --git a/tests/test_utils/test_auth.py b/tests/test_utils/test_auth.py index 4465f9a..a13a4b9 100644 --- a/tests/test_utils/test_auth.py +++ b/tests/test_utils/test_auth.py @@ -1,6 +1,10 @@ +from unittest.mock import patch + import pytest -from favorites_crawler.utils.auth import parse_twitter_likes_url, parser_twitter_likes_features +from favorites_crawler.utils.auth import parse_twitter_likes_url, parser_twitter_likes_features, \ + refresh_pixiv_token +from favorites_crawler.utils.config import dump_config, load_config def test_parser_twitter_likes_features(): @@ -36,3 +40,35 @@ def test_parser_twitter_likes_features(): )) def test_twitter_parse_likes_url(url, expected): assert parse_twitter_likes_url(url) == expected + + +@patch('favorites_crawler.utils.auth.CustomGetPixivToken') +class TestRefreshPixivToken: + def test_should_return_access_token_if_refresh_token_exists(self, mock_gppt, tmp_path): + favors_home = tmp_path / 'home' + refresh_token = 'refresh_token' + new_access_token = 'new_access_token' + config = { + 'pixiv': { + 'ACCESS_TOKEN': 'old_access_token', + 'REFRESH_TOKEN': refresh_token, + } + } + dump_config(config, favors_home) + mock_refresh = mock_gppt.return_value.refresh + mock_refresh.return_value = {'access_token': new_access_token} + + access_token = refresh_pixiv_token(favors_home) + + mock_refresh.assert_called_once_with(refresh_token) + assert access_token == new_access_token + assert load_config(favors_home)['pixiv']['ACCESS_TOKEN'] == new_access_token + + def test_should_raise_value_error_if_refresh_token_not_exists(self, mock_gppt, tmp_path): + favors_home = tmp_path / 'home' + mock_refresh = mock_gppt.return_value.refresh + + with pytest.raises(ValueError): + refresh_pixiv_token(favors_home) + + mock_refresh.assert_not_called()