holy-python/tests/test_acceptance_creation.py

import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from acceptance_service import (
    XHS_URL_PATTERN,
    cancel_acceptance_task,
    collect_acceptance_now,
    collect_acceptance_crawler_now,
    create_acceptance_task,
    get_acceptance_task,
    run_due_acceptance_schedules,
    retry_failed_acceptance_devices,
    stop_acceptance_schedule,
    update_acceptance_task,
    update_acceptance_schedule,
)
from config import Config
from database import connection, init_database, now_iso


class _ThreadWithoutStart:
    def __init__(self, *args, **kwargs):
        pass

    def start(self):
        pass


class AcceptanceCreationTest(unittest.TestCase):
    def test_xhs_short_link_cn_is_recognized(self):
        text = "笔记 https://xhslink.cn/o/example123 进入小红书"
        self.assertEqual(XHS_URL_PATTERN.findall(text), ["https://xhslink.cn/o/example123"])

    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.original_database_path = Config.DATABASE_PATH
        self.original_file_storage_path = Config.FILE_STORAGE_PATH
        root = Path(self.temp_dir.name)
        Config.DATABASE_PATH = root / "holy.db"
        Config.FILE_STORAGE_PATH = root / "files"
        init_database()

    def tearDown(self):
        Config.DATABASE_PATH = self.original_database_path
        Config.FILE_STORAGE_PATH = self.original_file_storage_path
        self.temp_dir.cleanup()

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.get_wen_resources")
    def test_screenshot_task_requires_keywords_but_not_feishu_links(self, resources):
        resources.return_value = {"cookie_count": 0, "device_count": 2}

        task = create_acceptance_task({
            "acceptance_method": "screenshot",
            "feishu_links": [],
            "keywords": ["山姆特价"],
            "device_count": 1,
        })

        self.assertEqual(task["acceptance_method"], "screenshot")
        self.assertEqual(task["keywords"], ["山姆特价"])
        self.assertEqual(task["feishu_links"], [])

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.get_wen_resources")
    def test_create_can_save_schedule_without_starting_initial_collection(self, resources):
        resources.return_value = {"cookie_count": 2, "device_count": 3, "credential_ids": ["c1", "c2"]}
        task = create_acceptance_task({
            "acceptance_method": "combined", "keywords": ["山姆会员卡"],
            "match_keywords": ["民生"], "schedule_weekdays": [1, 4],
            "schedule_start_time": "20:00", "schedule_end_time": "22:00",
            "defer_initial_collection": True,
        })
        self.assertEqual(task["status"], "completed")
        self.assertEqual(task["cookie_count"], 2)
        self.assertEqual(task["device_count"], 3)
        self.assertEqual(task["schedule_weekdays"], [1, 4])
        self.assertEqual(task["schedule_start_time"], "20:00")
        self.assertIsNotNone(task["next_run_at"])
        self.assertEqual(task["keyword_runs"], [])

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.sync_dashboard_acceptance_project")
    @patch("acceptance_service.get_wen_resources")
    def test_project_update_can_clear_feishu_links_and_save_schedule(self, resources, sync_dashboard):
        resources.return_value = {"cookie_count": 2, "device_count": 3, "credential_ids": ["c1", "c2"]}
        task = create_acceptance_task({
            "acceptance_method": "combined", "name": "原项目",
            "feishu_links": ["https://example.feishu.cn/docx/project"],
            "keywords": ["山姆会员卡"], "match_keywords": ["民生"],
            "schedule_weekdays": [1], "schedule_start_time": "20:00",
            "schedule_end_time": "21:00", "defer_initial_collection": True,
        })

        updated = update_acceptance_task(task["task_id"], {
            "name": "", "feishu_links": [], "keywords": ["山姆回馈金"],
            "match_keywords": ["民生山姆卡"], "schedule_weekdays": [2, 5],
            "schedule_start_time": "19:30", "schedule_end_time": "22:00",
        })

        self.assertEqual(updated["name"], "项目验收｜山姆回馈金")
        self.assertEqual(updated["feishu_links"], [])
        self.assertEqual(updated["schedule_weekdays"], [2, 5])
        self.assertEqual(updated["schedule_start_time"], "19:30")
        self.assertEqual(updated["schedule_end_time"], "22:00")
        sync_dashboard.assert_called_once()

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.get_wen_resources")
    def test_active_acceptance_task_can_be_cancelled(self, resources):
        resources.return_value = {"cookie_count": 0, "device_count": 2}
        task = create_acceptance_task({
            "acceptance_method": "screenshot",
            "keywords": ["山姆特价"],
            "device_count": 1,
        })
        update_acceptance_schedule(task["task_id"], [1], "00:00", "23:59")

        cancelled = cancel_acceptance_task(task["task_id"])

        self.assertEqual(cancelled["status"], "cancelled")
        self.assertTrue(cancelled["schedule_enabled"])
        self.assertIsNotNone(cancelled["next_run_at"])
        self.assertEqual(cancel_acceptance_task(task["task_id"])["status"], "cancelled")

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.get_wen_resources")
    def test_custom_cookie_is_stored_as_secret_and_not_returned(self, resources):
        resources.return_value = {"cookie_count": 0, "device_count": 0, "credential_ids": []}
        custom_cookie = "web_session=secret-value; a1=another-secret"

        task = create_acceptance_task({
            "acceptance_method": "data-analysis",
            "feishu_links": ["https://example.feishu.cn/docx/project"],
            "keywords": ["山姆特价"],
            "custom_cookie": custom_cookie,
        })

        self.assertEqual(task["device_count"], 1)
        self.assertNotIn(custom_cookie, str(task))
        self.assertTrue(task["resource_progress"][0]["resource_id"].startswith("custom-"))
        with connection() as conn:
            row = conn.execute(
                "SELECT secret_value FROM acceptance_task_secrets WHERE secret_type='custom_cookie'"
            ).fetchone()
        self.assertEqual(row["secret_value"], custom_cookie)

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.get_wen_resources")
    def test_weekly_schedule_and_immediate_collection(self, resources):
        resources.return_value = {"cookie_count": 0, "device_count": 2}
        source = create_acceptance_task({
            "acceptance_method": "screenshot",
            "keywords": ["山姆特价"],
            "device_count": 1,
        })

        scheduled = update_acceptance_schedule(
            source["task_id"], [1, 3, 5], "01:30", "04:30",
        )
        self.assertTrue(scheduled["schedule_enabled"])
        self.assertEqual(scheduled["schedule_weekday"], 1)
        self.assertEqual(scheduled["schedule_weekdays"], [1, 3, 5])
        self.assertEqual(scheduled["schedule_start_time"], "01:30")
        self.assertEqual(scheduled["schedule_end_time"], "04:30")
        self.assertIsNotNone(scheduled["next_run_at"])

        child = collect_acceptance_now(source["task_id"])
        self.assertNotEqual(child["task_id"], source["task_id"])
        self.assertEqual(child["keywords"], ["山姆特价"])
        self.assertFalse(child["schedule_enabled"])

        stopped = stop_acceptance_schedule(source["task_id"])
        self.assertFalse(stopped["schedule_enabled"])
        self.assertIsNone(stopped["next_run_at"])

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.get_wen_resources")
    def test_screenshot_project_can_start_cookie_crawler_batch(self, resources):
        resources.return_value = {"cookie_count": 2, "device_count": 3, "credential_ids": ["cookie-1", "cookie-2"]}
        source = create_acceptance_task({
            "acceptance_method": "screenshot", "keywords": ["山姆会员卡"],
            "match_keywords": ["民生"], "device_count": 3,
        })

        child = collect_acceptance_crawler_now(source["task_id"])

        self.assertEqual(child["acceptance_method"], "combined")
        self.assertEqual(child["keywords"], ["山姆会员卡"])
        self.assertEqual(child["match_keywords"], ["民生"])
        self.assertEqual(child["cookie_count"], 2)
        self.assertIsNotNone(child["batch_date"])

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service._copy_acceptance_task")
    @patch("acceptance_service.get_wen_resources")
    def test_due_schedule_reads_persisted_multi_day_range(self, resources, copy_task):
        resources.return_value = {"cookie_count": 0, "device_count": 2}
        source = create_acceptance_task({
            "acceptance_method": "screenshot",
            "keywords": ["山姆特价"],
            "device_count": 1,
        })
        update_acceptance_schedule(source["task_id"], [1, 4, 7], "00:00", "23:59")
        with connection() as conn:
            conn.execute(
                "UPDATE acceptance_tasks SET next_run_at=? WHERE public_task_id=?",
                (now_iso(), source["task_id"]),
            )

        self.assertEqual(run_due_acceptance_schedules(), 1)
        copy_task.assert_called_once_with(source["task_id"], trigger_source="acceptance-weekly")
        refreshed = get_acceptance_task(source["task_id"])
        self.assertEqual(refreshed["schedule_weekdays"], [1, 4, 7])
        self.assertIsNotNone(refreshed["last_run_at"])
        self.assertIsNotNone(refreshed["next_run_at"])

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.schedule_xapi_screenshots")
    @patch("acceptance_service.list_xhs_wen_xapi_devices")
    @patch("acceptance_service.get_wen_resources")
    def test_retry_failed_devices_dispatches_only_failed_accounts(self, resources, devices, schedule):
        resources.return_value = {"cookie_count": 0, "device_count": 2}
        devices.return_value = [
            {"xhs_account_id": "failed-account"},
            {"xhs_account_id": "successful-account"},
        ]
        schedule.return_value = {
            "execution_mode": "immediate",
            "batches": [{
                "batch_id": "retry-batch-1",
                "items": [{"xhs_account_id": "failed-account"}],
            }],
        }
        source = create_acceptance_task({
            "acceptance_method": "screenshot",
            "keywords": ["山姆特价"],
            "device_count": 2,
        })
        with connection() as conn:
            task = conn.execute(
                "SELECT id FROM acceptance_tasks WHERE public_task_id=?", (source["task_id"],),
            ).fetchone()
            now = now_iso()
            run = conn.execute(
                """INSERT INTO acceptance_keyword_runs
                   (task_id,keyword,run_type,run_order,status,created_at,updated_at)
                   VALUES (?,'山姆特价','screenshot',0,'failed',?,?)""",
                (task["id"], now, now),
            )
            conn.executemany(
                """INSERT INTO acceptance_task_resources
                   (task_id,resource_type,resource_id,resource_order,created_at)
                   VALUES (?,'device',?,?,?)""",
                [
                    (task["id"], "failed-account", 0, now),
                    (task["id"], "successful-account", 1, now),
                ],
            )
            conn.executemany(
                """INSERT INTO acceptance_xapi_runs
                   (task_id,keyword_run_id,batch_id,xapi_task_id,keyword,xhs_account_id,
                    status,artifacts_json,feishu_sync_status,raw_json,created_at,updated_at)
                   VALUES (?,?,?,?,?,?,?,?,?,?,?,?)""",
                [
                    (task["id"], run.lastrowid, "old-1", "task-failed", "山姆特价", "failed-account", "failed", "[]", "failed", "{}", now, now),
                    (task["id"], run.lastrowid, "old-2", "task-success", "山姆特价", "successful-account", "completed", '[{"image_url":"https://example.test/ok.png"}]', "succeeded", "{}", now, now),
                ],
            )

        retried = retry_failed_acceptance_devices(source["task_id"])

        self.assertEqual(retried["retry_dispatched_device_count"], 1)
        self.assertEqual(retried["retry_dispatched_pair_count"], 1)
        schedule.assert_called_once_with(
            ["山姆特价"], trigger_source="acceptance",
            xhs_account_ids=["failed-account"],
        )
        with connection() as conn:
            keyword_run = conn.execute(
                "SELECT status,crawler_task_id FROM acceptance_keyword_runs WHERE id=?",
                (run.lastrowid,),
            ).fetchone()
        self.assertEqual(keyword_run["status"], "waiting")
        self.assertEqual(keyword_run["crawler_task_id"], "retry-batch-1")

    @patch("acceptance_service.threading.Thread", _ThreadWithoutStart)
    @patch("acceptance_service.get_wen_resources")
    def test_combined_task_uses_all_resources_and_keeps_match_keywords(self, resources):
        resources.return_value = {
            "cookie_count": 3, "device_count": 5,
            "credential_ids": ["cookie-1", "cookie-2", "cookie-3"],
        }
        task = create_acceptance_task({
            "acceptance_method": "combined",
            "feishu_links": ["https://example.feishu.cn/docx/project"],
            "keywords": ["山姆特价"],
            "match_keywords": ["民生", "联名卡"],
        })
        self.assertEqual(task["device_count"], 5)
        self.assertEqual(task["cookie_count"], 3)
        self.assertEqual(task["match_keywords"], ["民生", "联名卡"])
        self.assertEqual(
            len([item for item in task["resource_progress"] if item["resource_type"] == "cookie"]),
            3,
        )


if __name__ == "__main__":
    unittest.main()
    cancel_acceptance_task,