Merge pull request #66 from RyouMon/fix-crawl-pixiv

fix: fix crawl pixiv

Author: RyouMon

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')

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}')

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,8 +26,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 | Path):
+    config = load_config(home)
     pixiv_config = config.get('pixiv', {})
     refresh_token = pixiv_config.get('REFRESH_TOKEN')
     if not refresh_token:
@@ -33,7 +36,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
 
 

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)

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()