feat(07-03): implement login CLI subcommand
- Add login_command function with QR code login, cookie import, and status check - Add login subcommand with --import-cookie and --status options - Update CLI help text with login examples - Fix mock paths in tests to match actual import locations Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -11,6 +11,9 @@ KB Workflow - Knowledge Base Workflow Orchestrator
|
||||
使用方式:
|
||||
python kb_workflow.py process "https://www.bilibili.com/video/BVxxx"
|
||||
python kb_workflow.py process "https://..." -o /path/to/output
|
||||
python kb_workflow.py login # QR code login
|
||||
python kb_workflow.py login --import-cookie cookies.json
|
||||
python kb_workflow.py login --status
|
||||
"""
|
||||
|
||||
import argparse
|
||||
@@ -102,6 +105,68 @@ def process_command(
|
||||
}
|
||||
|
||||
|
||||
def login_command(
|
||||
import_cookie: Optional[str] = None,
|
||||
status: bool = False
|
||||
) -> Dict[str, Any]:
|
||||
"""Execute login command
|
||||
|
||||
Args:
|
||||
import_cookie: Path to JSON cookie file
|
||||
status: Check current cookie status
|
||||
|
||||
Returns:
|
||||
Dictionary with login result
|
||||
"""
|
||||
try:
|
||||
from kb_workflow.auth import ZhihuLogin
|
||||
from kb_workflow.auth.cookie_manager import CookieManager, parse_json_cookies
|
||||
from kb_workflow.pipeline.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
manager = CookieManager()
|
||||
|
||||
# Check Cookie status
|
||||
if status:
|
||||
current_cookie = settings.zhihu_cookie.get_secret_value()
|
||||
if not current_cookie:
|
||||
return {"valid": False, "reason": "No cookie found in ~/.kb_env"}
|
||||
|
||||
async def check_status():
|
||||
login = ZhihuLogin(intermediate_path=settings.intermediate_path)
|
||||
try:
|
||||
login.browser = login.browser or type('obj', (object,), {
|
||||
'browser': None,
|
||||
'page': None,
|
||||
'start': lambda cookie=None: None,
|
||||
'close': lambda: None
|
||||
})()
|
||||
login.browser.start(cookie=current_cookie)
|
||||
return await login.validate_cookie(current_cookie)
|
||||
finally:
|
||||
if login.browser and hasattr(login.browser, 'close'):
|
||||
login.browser.close()
|
||||
|
||||
result = asyncio.run(check_status())
|
||||
return result.model_dump()
|
||||
|
||||
# Import Cookie
|
||||
if import_cookie:
|
||||
cookie_str = parse_json_cookies(import_cookie)
|
||||
manager.save_cookie(cookie_str)
|
||||
return {"success": True, "message": f"Cookie imported from {import_cookie}"}
|
||||
|
||||
# QR code login (synchronous method)
|
||||
login = ZhihuLogin(intermediate_path=settings.intermediate_path)
|
||||
result = login.login_by_qrcode()
|
||||
if result.success and result.cookie_str:
|
||||
manager.save_cookie(result.cookie_str)
|
||||
return result.model_dump()
|
||||
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
|
||||
def main():
|
||||
"""CLI入口 - 用于命令行调用"""
|
||||
parser = argparse.ArgumentParser(
|
||||
@@ -117,6 +182,15 @@ def main():
|
||||
|
||||
# 强制重新执行(忽略已完成状态)
|
||||
python kb_workflow.py process "https://..." --no-skip
|
||||
|
||||
# 知乎登录 (二维码扫码)
|
||||
python kb_workflow.py login
|
||||
|
||||
# 导入 Cookie
|
||||
python kb_workflow.py login --import-cookie cookies.json
|
||||
|
||||
# 检查 Cookie 状态
|
||||
python kb_workflow.py login --status
|
||||
'''
|
||||
)
|
||||
|
||||
@@ -132,6 +206,13 @@ def main():
|
||||
process_parser.add_argument('--no-skip', action='store_false', dest='skip_completed',
|
||||
help='Re-run all steps')
|
||||
|
||||
# Login subcommand
|
||||
login_parser = subparsers.add_parser('login', help='Zhihu login authentication')
|
||||
login_parser.add_argument('--import-cookie', '-i', type=str,
|
||||
help='Import cookie from JSON file')
|
||||
login_parser.add_argument('--status', action='store_true',
|
||||
help='Check current cookie status')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Handle process subcommand
|
||||
@@ -142,6 +223,13 @@ def main():
|
||||
skip_completed=args.skip_completed
|
||||
)
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
# Handle login subcommand
|
||||
elif args.command == 'login':
|
||||
result = login_command(
|
||||
import_cookie=args.import_cookie,
|
||||
status=args.status
|
||||
)
|
||||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
@@ -57,13 +57,14 @@ class TestLoginCommand:
|
||||
self, mock_settings, mock_login_result_success
|
||||
):
|
||||
"""Test 1: login_default_mode - Default starts QR code login."""
|
||||
with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings), \
|
||||
patch('kb_workflow.kb_workflow.CookieManager') as MockManager, \
|
||||
patch('kb_workflow.kb_workflow.asyncio.run') as mock_run, \
|
||||
patch('kb_workflow.kb_workflow.ZhihuLogin') as MockLogin:
|
||||
with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings), \
|
||||
patch('kb_workflow.auth.cookie_manager.CookieManager') as MockManager, \
|
||||
patch('kb_workflow.auth.ZhihuLogin') as MockLogin:
|
||||
|
||||
# Setup mock
|
||||
mock_run.return_value = mock_login_result_success.model_dump()
|
||||
mock_login_instance = MagicMock()
|
||||
mock_login_instance.login_by_qrcode.return_value = mock_login_result_success
|
||||
MockLogin.return_value = mock_login_instance
|
||||
mock_manager = MagicMock()
|
||||
MockManager.return_value = mock_manager
|
||||
|
||||
@@ -83,9 +84,9 @@ class TestLoginCommand:
|
||||
{"name": "z_c0", "value": "imported_value", "domain": ".zhihu.com"}
|
||||
]))
|
||||
|
||||
with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings), \
|
||||
patch('kb_workflow.kb_workflow.CookieManager') as MockManager, \
|
||||
patch('kb_workflow.kb_workflow.parse_json_cookies') as mock_parse:
|
||||
with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings), \
|
||||
patch('kb_workflow.auth.cookie_manager.CookieManager') as MockManager, \
|
||||
patch('kb_workflow.auth.cookie_manager.parse_json_cookies') as mock_parse:
|
||||
|
||||
mock_parse.return_value = "z_c0=imported_value"
|
||||
mock_manager = MagicMock()
|
||||
@@ -101,7 +102,7 @@ class TestLoginCommand:
|
||||
|
||||
def test_login_status_no_cookie(self, mock_settings):
|
||||
"""Test 3a: login_status - --status with no cookie returns invalid."""
|
||||
with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings):
|
||||
with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings):
|
||||
|
||||
# Act
|
||||
result = login_command(status=True)
|
||||
@@ -114,14 +115,14 @@ class TestLoginCommand:
|
||||
self, mock_settings_with_cookie, mock_validation_result_valid
|
||||
):
|
||||
"""Test 3b: login_status - --status with valid cookie returns user info."""
|
||||
with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings_with_cookie), \
|
||||
with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings_with_cookie), \
|
||||
patch('kb_workflow.kb_workflow.asyncio.run') as mock_run, \
|
||||
patch('kb_workflow.kb_workflow.ZhihuLogin') as MockLogin:
|
||||
patch('kb_workflow.auth.ZhihuLogin') as MockLogin:
|
||||
|
||||
# Setup mock
|
||||
# Setup mock - asyncio.run returns the result object directly
|
||||
mock_login_instance = MagicMock()
|
||||
MockLogin.return_value = mock_login_instance
|
||||
mock_run.return_value = mock_validation_result_valid.model_dump()
|
||||
mock_run.return_value = mock_validation_result_valid # Return object, not dict
|
||||
|
||||
# Act
|
||||
result = login_command(status=True)
|
||||
@@ -135,12 +136,14 @@ class TestLoginCommand:
|
||||
self, mock_settings, mock_login_result_success
|
||||
):
|
||||
"""Test 4: login_integration - Successful login saves cookie."""
|
||||
with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings), \
|
||||
patch('kb_workflow.kb_workflow.CookieManager') as MockManager, \
|
||||
patch('kb_workflow.kb_workflow.asyncio.run') as mock_run:
|
||||
with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings), \
|
||||
patch('kb_workflow.auth.cookie_manager.CookieManager') as MockManager, \
|
||||
patch('kb_workflow.auth.ZhihuLogin') as MockLogin:
|
||||
|
||||
# Setup mock
|
||||
mock_run.return_value = mock_login_result_success.model_dump()
|
||||
mock_login_instance = MagicMock()
|
||||
mock_login_instance.login_by_qrcode.return_value = mock_login_result_success
|
||||
MockLogin.return_value = mock_login_instance
|
||||
mock_manager = MagicMock()
|
||||
MockManager.return_value = mock_manager
|
||||
|
||||
@@ -155,7 +158,7 @@ class TestLoginCommand:
|
||||
|
||||
def test_login_error_handling(self, mock_settings):
|
||||
"""Test 5: login handles exceptions gracefully."""
|
||||
with patch('kb_workflow.kb_workflow.get_settings', side_effect=Exception("Config error")):
|
||||
with patch('kb_workflow.pipeline.config.get_settings', side_effect=Exception("Config error")):
|
||||
|
||||
# Act
|
||||
result = login_command()
|
||||
|
||||
Reference in New Issue
Block a user