getDomain/domain-api/tests/test_juming_services.py

from __future__ import annotations

import json
import os
import tempfile
import unittest
from contextlib import contextmanager
from unittest.mock import patch

from requests.cookies import RequestsCookieJar

from app.services import juming_service, juming_task_service


class JumingServiceTests(unittest.TestCase):
    def test_crawl_juming_rejects_invalid_cookie(self) -> None:
        cookie_jar = RequestsCookieJar()
        cookie_jar.set("sid", "expired")

        with patch.object(juming_service, "_load_juming_cookie", return_value=(cookie_jar, "local")):
            with patch.object(juming_service, "_validate_juming_cookie", return_value=(False, "聚名登录态已失效，请重新登录")):
                with self.assertRaisesRegex(ValueError, "已失效"):
                    juming_service.crawl_juming({"mode": "delete_list"})

    def test_get_juming_status_exposes_remote_validation(self) -> None:
        cookie_jar = RequestsCookieJar()
        cookie_jar.set("sid", "alive")

        with patch.object(juming_service, "_load_juming_cookie", return_value=(cookie_jar, "local")):
            with patch.object(juming_service, "_validate_juming_cookie", return_value=(False, "聚名登录态已失效，请重新登录")):
                payload = juming_service.get_juming_status()

        self.assertTrue(payload["cookie_present"])
        self.assertFalse(payload["cookie_valid"])
        self.assertFalse(payload["cookie_ready"])
        self.assertIn("已失效", payload["cookie_message"])

    def test_insert_domains_uses_copy_stage_import_path(self) -> None:
        class FakeCursor:
            def __init__(self) -> None:
                self.executed: list[tuple[str, object]] = []
                self.copy_calls: list[tuple[str, tuple[str, ...], str]] = []

            def __enter__(self):
                return self

            def __exit__(self, exc_type, exc, tb):
                return False

            def execute(self, sql: str, params: object = None) -> None:
                self.executed.append((sql, params))

            def copy_from(self, file_obj, table: str, columns: tuple[str, ...]) -> None:
                self.copy_calls.append((table, columns, file_obj.read()))

            def fetchone(self):
                return (2, 2, 0)

        class FakeConn:
            def __init__(self) -> None:
                self.cursor_obj = FakeCursor()
                self.commit_calls = 0

            def cursor(self):
                return self.cursor_obj

            def commit(self) -> None:
                self.commit_calls += 1

        fake_conn = FakeConn()

        @contextmanager
        def fake_get_db():
            yield fake_conn

        with patch.object(juming_service, "get_db", fake_get_db):
            with patch.object(juming_service, "IMPORT_BATCH_SIZE", 2):
                stats = juming_service._insert_domains(
                    ["alpha.com", "beta.net"],
                    juming_service.DELETE_LIST_SOURCE_TYPE,
                )

        self.assertEqual(2, stats["added"])
        self.assertEqual(1, fake_conn.commit_calls)
        self.assertTrue(fake_conn.cursor_obj.copy_calls)
        table, columns, payload = fake_conn.cursor_obj.copy_calls[0]
        self.assertEqual("juming_import_stage", table)
        self.assertEqual(("domain", "tld"), columns)
        self.assertIn("alpha.com\tcom", payload)
        self.assertIn("beta.net\tnet", payload)
        executed_sql = "\n".join(sql for sql, _params in fake_conn.cursor_obj.executed)
        self.assertIn("create temporary table if not exists juming_import_stage", executed_sql.lower())
        self.assertIn("inserted as", executed_sql.lower())
        self.assertIn("left join domains existing", executed_sql.lower())

    def test_delete_list_import_skips_already_imported_same_signature(self) -> None:
        cookie_jar = RequestsCookieJar()
        cookie_jar.set("sid", "alive")

        class FakeJM:
            cookie = cookie_jar

            def new_cha_del(self, current_date: str):
                if current_date == "2026-03-21":
                    return ["alpha.com", "beta.net"]
                return []

        logs: list[str] = []
        signature = juming_service._compute_domains_signature(["alpha.com", "beta.net"])
        cached_state = {
            "2026-03-21": {
                "signature": signature,
                "total": 2,
                "valid": 2,
                "invalid": 0,
            }
        }

        with patch.object(juming_service, "_load_juming_cookie", return_value=(cookie_jar, "local")):
            with patch.object(juming_service, "JM", return_value=FakeJM()):
                with patch.object(juming_service, "_load_delete_import_state", return_value=cached_state):
                    with patch.object(juming_service, "_insert_domains") as mock_insert:
                        result = juming_service._crawl_delete_list_and_import(
                            "2026-03-21",
                            False,
                            log=logs.append,
                        )

        mock_insert.assert_not_called()
        self.assertEqual(0, result["stats"]["added"])
        self.assertEqual(2, result["stats"]["exists"])
        self.assertTrue(any("跳过重复入库" in line for line in logs))


class JumingTaskServiceTests(unittest.TestCase):
    def setUp(self) -> None:
        self._tmpdir = tempfile.TemporaryDirectory()
        self._old_runtime_root = os.environ.get("DOMAIN_API_RUNTIME_ROOT")
        os.environ["DOMAIN_API_RUNTIME_ROOT"] = self._tmpdir.name
        juming_task_service._ACTIVE_TASK_IDS.clear()

    def tearDown(self) -> None:
        juming_task_service._ACTIVE_TASK_IDS.clear()
        if self._old_runtime_root is None:
            os.environ.pop("DOMAIN_API_RUNTIME_ROOT", None)
        else:
            os.environ["DOMAIN_API_RUNTIME_ROOT"] = self._old_runtime_root
        self._tmpdir.cleanup()

    def _write_tasks(self, records: list[dict]) -> None:
        path = os.path.join(self._tmpdir.name, "juming_tasks.json")
        with open(path, "w", encoding="utf-8") as handle:
            json.dump(records, handle, ensure_ascii=False, indent=2)

    def test_cleanup_orphaned_tasks_marks_running_task_failed(self) -> None:
        self._write_tasks(
            [
                {
                    "task_id": "task-1",
                    "status": "running",
                    "phase": "importing",
                    "phase_label": "入库中",
                    "cancel_requested": False,
                    "message": "开始入库处理",
                    "created_at": "2026-04-21 21:00:00",
                    "updated_at": "2026-04-21 21:00:00",
                    "started_at": "2026-04-21 21:00:00",
                    "completed_at": "",
                    "result": None,
                    "logs": [],
                }
            ]
        )

        tasks = juming_task_service.list_juming_tasks()
        self.assertEqual("failed", tasks[0]["status"])
        self.assertIn("中断", tasks[0]["message"])

    def test_create_task_rejects_parallel_active_task(self) -> None:
        self._write_tasks(
            [
                {
                    "task_id": "task-1",
                    "status": "running",
                    "phase": "fetching",
                    "phase_label": "抓取中",
                    "cancel_requested": False,
                    "message": "正在抓取",
                    "created_at": "2026-04-21 21:00:00",
                    "updated_at": "2026-04-21 21:00:00",
                    "started_at": "2026-04-21 21:00:00",
                    "completed_at": "",
                    "result": None,
                    "logs": [],
                }
            ]
        )
        juming_task_service._ACTIVE_TASK_IDS.add("task-1")

        with self.assertRaisesRegex(ValueError, "已有聚名采集任务正在运行"):
            juming_task_service.create_juming_task({"mode": "delete_list"})


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