diff --git a/Douyin.py b/Douyin.py index d4073fa..c5b78eb 100644 --- a/Douyin.py +++ b/Douyin.py @@ -547,6 +547,12 @@ def build_parser() -> argparse.ArgumentParser: default=DEFAULT_BROWSER_PORT, help="附着到已启动 Chrome 的调试端口,默认 9223", ) + parser.add_argument( + "--max-videos", + type=int, + default=50, + help="推荐流最大抓取数量,默认 50", + ) return parser @@ -560,6 +566,8 @@ def main(argv: list[str] | None = None) -> int: parser.error("--timeout 必须大于 0") if args.browser_port is not None and args.browser_port <= 0: parser.error("--browser-port 必须大于 0") + if args.max_videos <= 0: + parser.error("--max-videos 必须大于 0") try: target = resolve_cli_target(args.target, browser_port=args.browser_port) @@ -572,6 +580,13 @@ def main(argv: list[str] | None = None) -> int: browser_port=args.browser_port, auto_scroll=args.pages > 1, ) + elif target.kind == "recommendation": + total = collect_recommendations( + max_videos=args.max_videos, + timeout=args.timeout, + output_dir=Path(args.output_dir), + browser_port=args.browser_port, + ) elif target.kind == "single-video": total = collect_single_video( target=target, diff --git a/test_douyin.py b/test_douyin.py index 3333bf3..99d9dcb 100644 --- a/test_douyin.py +++ b/test_douyin.py @@ -404,6 +404,11 @@ class DouyinModuleTests(unittest.TestCase): self.assertEqual(page.visited_urls, ["https://www.douyin.com/video/7619989983668240802"]) mocked_download.assert_called_once() + def test_build_parser_has_max_videos_argument(self) -> None: + module = importlib.import_module("Douyin") + args = module.build_parser().parse_args(["--max-videos", "30"]) + self.assertEqual(args.max_videos, 30) + def test_build_parser_defaults_to_zero_argument_current_page_flow(self) -> None: module = importlib.import_module("Douyin") args = module.build_parser().parse_args([]) @@ -419,6 +424,26 @@ class DouyinModuleTests(unittest.TestCase): self.assertEqual(target.aweme_id, "7619989983668240802") mocked_imports.assert_not_called() + def test_main_dispatches_recommendation_flow_for_recommendation_url(self) -> None: + module = importlib.import_module("Douyin") + stdout = io.StringIO() + recommendation_target = module.ResolvedTarget( + kind="recommendation", + value="https://www.douyin.com/", + source="current-page", + ) + with redirect_stdout(stdout): + with mock.patch.object(module, "resolve_cli_target", return_value=recommendation_target): + with mock.patch.object(module, "collect_recommendations", return_value=5) as mocked_collect: + exit_code = module.main([]) + self.assertEqual(exit_code, 0) + mocked_collect.assert_called_once_with( + max_videos=50, + timeout=10, + output_dir=module.Path("video"), + browser_port=9223, + ) + def test_main_without_target_dispatches_current_page_creator_flow(self) -> None: module = importlib.import_module("Douyin") stdout = io.StringIO()