迁移目录
This commit is contained in:
hajimi
@@ -0,0 +1,203 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user