204 lines
7.6 KiB
Python
204 lines
7.6 KiB
Python
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
import importlib
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
|
|
class FakeStore:
|
|
def __init__(self):
|
|
self.created = []
|
|
self.progress = []
|
|
self.finished = []
|
|
self.next_id = 100
|
|
|
|
def create_running(self, task):
|
|
self.created.append(task)
|
|
self.next_id += 1
|
|
return self.next_id
|
|
|
|
def update_progress(self, log_id, output, elapsed):
|
|
self.progress.append((log_id, output, elapsed))
|
|
|
|
def finish(self, log_id, status, output, elapsed, error_message=""):
|
|
self.finished.append((log_id, status, output, elapsed, error_message))
|
|
|
|
|
|
class DockerTaskRunnerTest(unittest.TestCase):
|
|
def test_config_uses_dqd_environment_overrides(self):
|
|
sys.modules.pop("src.core.config", None)
|
|
src_core = importlib.import_module("src.core")
|
|
if hasattr(src_core, "config"):
|
|
delattr(src_core, "config")
|
|
config_module = importlib.import_module("src.core.config")
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
settings = Path(tmp) / "settings.yaml"
|
|
settings.write_text(
|
|
"\n".join(
|
|
[
|
|
"database:",
|
|
' host: "yaml-db"',
|
|
" port: 3306",
|
|
' database: "yaml_name"',
|
|
' username: "yaml_user"',
|
|
' password: "yaml_password"',
|
|
' prefix: "la_"',
|
|
"redis:",
|
|
' host: "yaml-redis"',
|
|
" port: 6379",
|
|
" db: 0",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
env = {
|
|
"DQD_DB_HOST": "docker-db",
|
|
"DQD_DB_PORT": "3307",
|
|
"DQD_DB_NAME": "docker_name",
|
|
"DQD_DB_USER": "docker_user",
|
|
"DQD_DB_PASSWORD": "docker_password",
|
|
"DQD_DB_PREFIX": "sx_",
|
|
"DQD_REDIS_HOST": "docker-redis",
|
|
"DQD_REDIS_PORT": "6380",
|
|
"DQD_REDIS_DB": "5",
|
|
}
|
|
with patch.dict(os.environ, env, clear=False):
|
|
config_module._config = None
|
|
cfg = config_module.load_config(str(settings))
|
|
|
|
self.assertEqual(cfg.database.host, "docker-db")
|
|
self.assertEqual(cfg.database.port, 3307)
|
|
self.assertEqual(cfg.database.database, "docker_name")
|
|
self.assertEqual(cfg.database.username, "docker_user")
|
|
self.assertEqual(cfg.database.password, "docker_password")
|
|
self.assertEqual(cfg.database.prefix, "sx_")
|
|
self.assertEqual(cfg.redis.host, "docker-redis")
|
|
self.assertEqual(cfg.redis.port, 6380)
|
|
self.assertEqual(cfg.redis.db, 5)
|
|
|
|
def test_render_crontab_skips_inactive_tasks_and_quotes_arguments(self):
|
|
from scripts.docker_task_runner import CrawlerTask, render_crontab
|
|
|
|
tasks = [
|
|
CrawlerTask("nba", "NBA赛程", "nba_match", "*/1 * * * *", True),
|
|
CrawlerTask("old", "旧任务", "standings", "0 2 * * *", False),
|
|
]
|
|
|
|
rendered = render_crontab(tasks, python_bin="/usr/local/bin/python", app_dir="/app")
|
|
|
|
self.assertIn("SHELL=/bin/bash", rendered)
|
|
self.assertIn("BASH_ENV=/etc/sport-era-crawler/env.sh", rendered)
|
|
self.assertIn("*/1 * * * * root cd /app && /usr/local/bin/python scripts/docker_task_runner.py run nba_match --task-key nba --task-name", rendered)
|
|
self.assertIn("--cron-expression '*/1 * * * *'", rendered)
|
|
self.assertNotIn("standings", rendered)
|
|
|
|
def test_load_tasks_accepts_top_level_list_yaml(self):
|
|
from scripts.docker_task_runner import load_tasks
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tasks_file = Path(tmp) / "crawler_tasks.yaml"
|
|
tasks_file.write_text(
|
|
"\n".join(
|
|
[
|
|
"- key: live",
|
|
" name: 实时详情",
|
|
" action: live_detail",
|
|
" cron: '*/1 * * * *'",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
tasks = load_tasks(tasks_file)
|
|
|
|
self.assertEqual(len(tasks), 1)
|
|
self.assertEqual(tasks[0].task_key, "live")
|
|
self.assertEqual(tasks[0].action, "live_detail")
|
|
|
|
def test_run_task_marks_success_and_flushes_output(self):
|
|
from scripts.docker_task_runner import CommandResult, CrawlerTask, run_task
|
|
|
|
store = FakeStore()
|
|
task = CrawlerTask("nba", "NBA赛程", "nba_match", "*/1 * * * *", True, timeout=3)
|
|
|
|
def fake_runner(action, timeout, output_callback):
|
|
output_callback("hello\n")
|
|
return CommandResult(exit_code=0, output="hello\nok\n", elapsed=0.2, timed_out=False)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
result = run_task(task, store=store, lock_dir=Path(tmp), command_runner=fake_runner)
|
|
|
|
self.assertEqual(result.status, 1)
|
|
self.assertEqual(store.created[0].action, "nba_match")
|
|
self.assertEqual(store.progress[0][1], "hello\n")
|
|
self.assertEqual(store.finished[0][1], 1)
|
|
self.assertIn("ok", store.finished[0][2])
|
|
|
|
def test_run_task_marks_failure(self):
|
|
from scripts.docker_task_runner import CommandResult, CrawlerTask, run_task
|
|
|
|
store = FakeStore()
|
|
task = CrawlerTask("news", "资讯", "league_news", "*/30 * * * *", True, timeout=3)
|
|
|
|
def fake_runner(action, timeout, output_callback):
|
|
return CommandResult(exit_code=7, output="boom\n", elapsed=0.1, timed_out=False)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
result = run_task(task, store=store, lock_dir=Path(tmp), command_runner=fake_runner)
|
|
|
|
self.assertEqual(result.status, 2)
|
|
self.assertEqual(store.finished[0][1], 2)
|
|
self.assertIn("退出码 7", store.finished[0][4])
|
|
|
|
def test_run_task_marks_timeout_as_failure(self):
|
|
from scripts.docker_task_runner import CommandResult, CrawlerTask, run_task
|
|
|
|
store = FakeStore()
|
|
task = CrawlerTask("live", "实时", "live_detail", "*/1 * * * *", True, timeout=1)
|
|
|
|
def fake_runner(action, timeout, output_callback):
|
|
return CommandResult(exit_code=-1, output="slow\n", elapsed=1.0, timed_out=True)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
result = run_task(task, store=store, lock_dir=Path(tmp), command_runner=fake_runner)
|
|
|
|
self.assertEqual(result.status, 2)
|
|
self.assertIn("执行超时", store.finished[0][4])
|
|
|
|
def test_run_task_skips_duplicate_action(self):
|
|
from scripts.docker_task_runner import CrawlerTask, TaskLock, run_task
|
|
|
|
store = FakeStore()
|
|
task = CrawlerTask("live", "实时", "live_detail", "*/1 * * * *", True, timeout=1)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
lock_dir = Path(tmp)
|
|
lock = TaskLock("live_detail", lock_dir)
|
|
self.assertTrue(lock.acquire())
|
|
try:
|
|
result = run_task(
|
|
task,
|
|
store=store,
|
|
lock_dir=lock_dir,
|
|
command_runner=lambda action, timeout, output_callback: self.fail("duplicate should not run"),
|
|
)
|
|
finally:
|
|
lock.release()
|
|
|
|
self.assertEqual(result.status, 3)
|
|
self.assertEqual(store.finished[0][1], 3)
|
|
self.assertIn("已有任务执行中", store.finished[0][2])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|