vmess-domain-rotator/tests/test_cfst_local_lifecycle.py

import csv
import json
import os
import pathlib
import shutil
import subprocess
import sys
import tempfile
import unittest

PROJECT_ROOT = pathlib.Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT / "scripts"))

from rotator.main import run_updater


class TestCfstLocalLifecycle(unittest.TestCase):
    """Local Python cfst_local 模式生命周期与选优逻辑测试套件"""

    def setUp(self):
        self.tmp_dir = tempfile.TemporaryDirectory()
        self.work_dir = pathlib.Path(self.tmp_dir.name)

    def tearDown(self):
        self.tmp_dir.cleanup()

    def _create_mock_cfst_env(self, csv_rows: list, run_args=None, skip_run=True, last_good=None) -> tuple[pathlib.Path, pathlib.Path]:
        """创建隔离的 cfst_local 运行目录、配置文件与模拟 CSV 数据"""
        cfst_dir = self.work_dir / "cfst_test"
        cfst_dir.mkdir(parents=True, exist_ok=True)
        csv_path = cfst_dir / "result.csv"

        with open(csv_path, "w", encoding="utf-8", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["IP 地址", "已发送", "已接收", "丢包率", "平均延迟", "下载速度 (MB/s)", "地区码"])
            for row in csv_rows:
                writer.writerow(row)

        runtime_dir = self.work_dir / "runtime"
        runtime_dir.mkdir(parents=True, exist_ok=True)
        if last_good:
            (runtime_dir / "state.json").write_text(json.dumps({
                "last_good_ip": last_good,
                "updated_at": "2026-06-01T00:00:00Z"
            }))

        config = {
            "source": {"type": "cfst_local"},
            "cfst_local": {
                "work_dir": str(cfst_dir),
                "binary": "./cfst",
                "run_args": run_args or ["-f", "ip.txt", "-o", "result.csv", "-p", "10"],
                "result_file": "result.csv",
                "run_timeout_sec": 60,
                "skip_run": skip_run,
                "header_rows": 1,
                "columns": {
                    "ip": 0,
                    "sent": 1,
                    "received": 2,
                    "loss_rate": 3,
                    "avg_latency": 4,
                    "download_speed": 5,
                    "region": 6,
                },
            },
            "domain_filter": {
                "include_suffixes": [],
                "exclude_regex": [],
            },
            "selection": {"top_n": 3},
            "output": {
                "runtime_dir": str(runtime_dir),
                "selected_value_file": "current_ip.txt",
                "selected_value_json": "current_ip.json",
                "selected_value_json_key": "ip",
                "state_file": "state.json",
                "state_last_good_key": "last_good_ip",
                "export_vars_file": "substore_vars.json",
                "substore_value_key": "AUTO_CFIP",
            },
            "notify": {"command": ""},
        }

        config_file = self.work_dir / "test_local.json"
        config_file.write_text(json.dumps(config, indent=2))
        return config_file, runtime_dir

    def test_cfst_local_standard_run_and_output_artifacts(self):
        """[全流程 1] 标准 cfst_local 流程：解析 CSV -> 过滤选优 -> 4 个 runtime 文件落盘 -> Sub-Store 契约输出"""
        rows = [
            ["172.64.147.31", "4", "4", "0.00", "22.14", "15.30", "HKG"],
            ["198.41.209.213", "4", "4", "0.00", "21.87", "12.80", "HKG"],
            ["104.19.38.34", "4", "4", "0.00", "22.82", "8.50", "HKG"],
            ["104.16.254.215", "4", "4", "0.00", "23.29", "5.20", "HKG"],
        ]
        config_file, runtime_dir = self._create_mock_cfst_env(rows)

        rc = run_updater(str(config_file))
        self.assertEqual(rc, 0)

        # 验证 4 份落盘文件
        txt_path = runtime_dir / "current_ip.txt"
        json_path = runtime_dir / "current_ip.json"
        vars_path = runtime_dir / "substore_vars.json"
        state_path = runtime_dir / "state.json"

        self.assertTrue(txt_path.exists())
        self.assertTrue(json_path.exists())
        self.assertTrue(vars_path.exists())
        self.assertTrue(state_path.exists())

        ip_json = json.loads(json_path.read_text())
        self.assertEqual(ip_json["status"], "ok")
        self.assertEqual(ip_json["ip"], "172.64.147.31")
        self.assertEqual(ip_json["source_type"], "cfst_local")
        self.assertEqual(ip_json["source_count"], 4)

        self.assertEqual(txt_path.read_text().strip(), "172.64.147.31")

        ip_json = json.loads(json_path.read_text())
        self.assertEqual(ip_json["ip"], "172.64.147.31")
        self.assertEqual(ip_json["status"], "ok")
        self.assertEqual(len(ip_json["top_candidates"]), 3)
        self.assertEqual(ip_json["top_candidates"][0]["download_speed"], "15.30")
        self.assertEqual(ip_json["top_candidates"][0]["region"], "HKG")

        vars_json = json.loads(vars_path.read_text())
        self.assertEqual(vars_json["AUTO_CFIP"], "172.64.147.31")
        self.assertEqual(vars_json["STATUS"], "ok")

        state_json = json.loads(state_path.read_text())
        self.assertEqual(state_json["last_good_ip"], "172.64.147.31")
        self.assertEqual(state_json["status"], "ok")

    def test_cfst_local_filters_zero_download_speed_and_prefers_positive(self):
        """[边界 2] 0 测速过滤：前两个节点测速为 0.00，自动跳过并选出首个 download_speed > 0 的健康节点"""
        rows = [
            ["1.1.1.1", "4", "4", "0.00", "15.00", "0.00", "N/A"],
            ["2.2.2.2", "4", "4", "0.00", "18.00", "0.00", "N/A"],
            ["3.3.3.3", "4", "4", "0.00", "25.00", "9.80", "HKG"],
            ["4.4.4.4", "4", "4", "0.00", "28.00", "6.20", "SJC"],
        ]
        config_file, runtime_dir = self._create_mock_cfst_env(rows)

        rc = run_updater(str(config_file))
        self.assertEqual(rc, 0)
        res = json.loads((runtime_dir / "current_ip.json").read_text())
        self.assertEqual(res["status"], "ok")
        self.assertEqual(res["ip"], "3.3.3.3", "必须自动过滤 0 速度死节点，挑选 3.3.3.3")

        candidates = [c["ip"] for c in res["top_candidates"]]
        self.assertNotIn("1.1.1.1", candidates, "0 测速节点不应出现在候选列表中")
        self.assertNotIn("2.2.2.2", candidates, "0 测速节点不应出现在候选列表中")
        self.assertIn("3.3.3.3", candidates)
        self.assertIn("4.4.4.4", candidates)

    def test_cfst_local_all_zero_speed_triggers_last_good_fallback(self):
        """[容灾 3] 全员测速 0.00 容灾回退：所有节点均不可用时，平滑回退至 state.json 中的 last_good_ip"""
        rows = [
            ["1.1.1.1", "4", "4", "0.00", "15.00", "0.00", "N/A"],
            ["2.2.2.2", "4", "4", "0.00", "18.00", "0.00", "N/A"],
        ]
        config_file, runtime_dir = self._create_mock_cfst_env(rows, last_good="198.51.100.88")

        rc = run_updater(str(config_file))
        self.assertEqual(rc, 0)
        res = json.loads((runtime_dir / "current_ip.json").read_text())
        self.assertEqual(res["status"], "fallback_last_good")
        self.assertEqual(res["ip"], "198.51.100.88")

        # 检查落盘文件
        txt_path = runtime_dir / "current_ip.txt"
        self.assertEqual(txt_path.read_text().strip(), "198.51.100.88")
        vars_path = runtime_dir / "substore_vars.json"
        vars_json = json.loads(vars_path.read_text())
        self.assertEqual(vars_json["AUTO_CFIP"], "198.51.100.88")
        self.assertEqual(vars_json["STATUS"], "fallback_last_good")

    def test_cfst_local_disable_download_flag_dd(self):
        """[参数 4] -dd 禁用下载测速模式：配置 -dd 时，按延迟排序，不因 download_speed == 0.00 误杀"""
        rows = [
            ["10.0.0.1", "4", "4", "0.00", "12.00", "0.00", "N/A"],
            ["10.0.0.2", "4", "4", "0.00", "15.00", "0.00", "N/A"],
        ]
        # 在 run_args 中注入 -dd
        config_file, runtime_dir = self._create_mock_cfst_env(rows, run_args=["-f", "ip.txt", "-dd", "-o", "result.csv"])

        rc = run_updater(str(config_file))
        self.assertEqual(rc, 0)
        res = json.loads((runtime_dir / "current_ip.json").read_text())
        self.assertEqual(res["status"], "ok")
        self.assertEqual(res["ip"], "10.0.0.1", "禁用下载测速模式下，按延迟排序取第一名")

    def test_cfst_local_filter_exclude_regex(self):
        """[过滤 5] 正则过滤排除指定 IP 段"""
        rows = [
            ["172.64.10.1", "4", "4", "0.00", "20.00", "20.00", "HKG"],
            ["104.16.10.1", "4", "4", "0.00", "22.00", "18.00", "HKG"],
        ]
        config_file, runtime_dir = self._create_mock_cfst_env(rows)

        # 增加正则过滤：排除 172.64.*
        with open(config_file, "r", encoding="utf-8") as f:
            cfg = json.load(f)
        cfg["domain_filter"]["exclude_regex"] = [r"^172\.64\."]
        with open(config_file, "w", encoding="utf-8") as f:
            json.dump(cfg, f, indent=2)

        rc = run_updater(str(config_file))
        self.assertEqual(rc, 0)
        res = json.loads((runtime_dir / "current_ip.json").read_text())
        self.assertEqual(res["status"], "ok")
        self.assertEqual(res["ip"], "104.16.10.1", "172.64.* 被排除，选取 104.16.10.1")

    def test_cfst_local_actual_repo_config_and_output_paths(self):
        """[契约 6] 验证仓库实际 configs/config.local.json 的输出路径解析与落盘契约"""
        real_cfg_path = PROJECT_ROOT / "configs" / "config.local.json"
        self.assertTrue(real_cfg_path.exists())

        cmd = [sys.executable, str(PROJECT_ROOT / "scripts" / "domain_updater.py"), "--config", str(real_cfg_path), "--print-output-settings"]
        out = subprocess.check_output(cmd, text=True).strip()
        settings = json.loads(out)

        # 必须锚定在项目根目录 /runtime，严禁出现 /configs/runtime
        expected_runtime_dir = str(PROJECT_ROOT / "runtime")
        self.assertEqual(settings["runtime_dir"], expected_runtime_dir)
        self.assertEqual(settings["selected_text_path"], str(PROJECT_ROOT / "runtime" / "current_ip.txt"))
        self.assertEqual(settings["selected_json_path"], str(PROJECT_ROOT / "runtime" / "current_ip.json"))
        self.assertEqual(settings["state_path"], str(PROJECT_ROOT / "runtime" / "state.json"))
        self.assertEqual(settings["vars_path"], str(PROJECT_ROOT / "runtime" / "substore_vars.json"))
        self.assertEqual(settings["selected_json_key"], "ip")
        self.assertEqual(settings["state_last_good_key"], "last_good_ip")
        self.assertEqual(settings["vars_value_key"], "AUTO_CFIP")


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