iAOP/deploy/orchestration/tests/test_rollback.py

# -*- coding: utf-8 -*-
"""部署回滚引擎测试（Issue #60）。

覆盖：
1. 快照压栈 + latest 查询；
2. FIFO 淘汰（超 max_snapshots）；
3. rollback_to_latest 成功（弹出栈顶 + 回滚目标版本）；
4. rollback(n) 多版本回溯；
5. 空栈回滚（NO_TARGET）；
6. restore 动作失败（FAILED + 压回栈顶）；
7. restore 异常（FAILED）；
8. post_rollback_health 失败（FAILED + 压回）；
9. n<1 跳过；
10. values_hash 漂移检测；
11. 报告序列化。
"""
import unittest

import _bootstrap  # noqa: F401  (sys.path 挂载)

from deploy_plan import DeployPlan
from rollback import (
    DeployError,
    RollbackManager,
    RollbackPoint,
    RollbackStatus,
)


def _restore(ok: bool, detail: str = ""):
    def _fn(snapshot, ctx):
        return ok, (detail or ("回滚成功" if ok else "回滚失败"))
    return _fn


def _raising_restore():
    def _fn(snapshot, ctx):
        raise RuntimeError("restore boom")
    return _fn


def _make_plan(version="v1.0.0", release="iaop"):
    return DeployPlan.default_helm_release(release=release, version=version)


class TestSnapshotStack(unittest.TestCase):
    """快照压栈与查询。"""

    def test_snapshot_push_and_latest(self):
        mgr = RollbackManager()
        # snapshot 即压栈（部署前快照），latest 即栈顶
        mgr.snapshot(_make_plan("v1.0.0"), values={"k": 1}, reason="第一次")
        self.assertEqual(mgr.latest().version, "v1.0.0")
        self.assertEqual(len(mgr.snapshots), 1)
        mgr.snapshot(_make_plan("v1.1.0"), values={"k": 2})
        self.assertEqual(mgr.latest().version, "v1.1.0")
        self.assertEqual(len(mgr.snapshots), 2)
        # 空栈时 latest 返回 None
        self.assertIsNotNone(mgr.latest())
        mgr2 = RollbackManager()
        self.assertIsNone(mgr2.latest())

    def test_fifo_eviction(self):
        mgr = RollbackManager(max_snapshots=2)
        mgr.snapshot(_make_plan("v1.0.0"))
        mgr.snapshot(_make_plan("v1.1.0"))
        mgr.snapshot(_make_plan("v1.2.0"))
        # 上限 2，淘汰最旧
        self.assertEqual(len(mgr.snapshots), 2)
        self.assertEqual([s.version for s in mgr.snapshots], ["v1.1.0", "v1.2.0"])

    def test_max_snapshots_validation(self):
        with self.assertRaises(DeployError):
            RollbackManager(max_snapshots=0)

    def test_values_hash_changes_with_config(self):
        mgr = RollbackManager()
        p1 = mgr.snapshot(_make_plan(), values={"backend": "gpu"})
        p2 = mgr.snapshot(_make_plan(), values={"backend": "npu"})
        self.assertNotEqual(p1.values_hash, p2.values_hash)


class TestRollbackSuccess(unittest.TestCase):
    """回滚成功场景。"""

    def test_rollback_to_latest_success(self):
        mgr = RollbackManager(restore_action=_restore(True))
        mgr.snapshot(_make_plan("v1.0.0"))
        mgr.snapshot(_make_plan("v1.1.0"))
        result = mgr.rollback_to_latest(reason="部署 v1.2.0 失败")
        self.assertTrue(result.succeeded)
        self.assertEqual(result.status, RollbackStatus.SUCCESS)
        self.assertEqual(result.target.version, "v1.1.0")
        # 弹出栈顶后栈剩 1 个
        self.assertEqual(len(mgr.snapshots), 1)
        self.assertEqual(mgr.latest().version, "v1.0.0")
        self.assertIn("v1.2.0 失败", result.reason)

    def test_rollback_n_versions(self):
        mgr = RollbackManager(restore_action=_restore(True))
        for v in ["v1.0.0", "v1.1.0", "v1.2.0"]:
            mgr.snapshot(_make_plan(v))
        # 回滚 2 个版本 → 目标 v1.1.0（弹出 v1.2.0 和 v1.1.0）
        result = mgr.rollback(n=2)
        self.assertTrue(result.succeeded)
        self.assertEqual(result.target.version, "v1.1.0")
        self.assertEqual(len(mgr.snapshots), 1)


class TestRollbackFailures(unittest.TestCase):
    """回滚失败与边界。"""

    def test_empty_stack_no_target(self):
        mgr = RollbackManager()
        result = mgr.rollback_to_latest()
        self.assertEqual(result.status, RollbackStatus.NO_TARGET)
        self.assertFalse(result.succeeded)

    def test_insufficient_stack_for_n(self):
        mgr = RollbackManager(restore_action=_restore(True))
        mgr.snapshot(_make_plan("v1.0.0"))
        # 栈仅 1 个，回滚 2 个 → NO_TARGET
        result = mgr.rollback(n=2)
        self.assertEqual(result.status, RollbackStatus.NO_TARGET)

    def test_restore_failure_pushes_target_back(self):
        mgr = RollbackManager(restore_action=_restore(False, "helm rollback 失败"))
        mgr.snapshot(_make_plan("v1.0.0"))
        mgr.snapshot(_make_plan("v1.1.0"))
        result = mgr.rollback_to_latest()
        self.assertEqual(result.status, RollbackStatus.FAILED)
        # 失败时目标压回栈顶，栈仍 2 个
        self.assertEqual(len(mgr.snapshots), 2)
        self.assertEqual(mgr.latest().version, "v1.1.0")

    def test_restore_exception_failure(self):
        mgr = RollbackManager(restore_action=_raising_restore())
        mgr.snapshot(_make_plan("v1.0.0"))
        result = mgr.rollback_to_latest()
        self.assertEqual(result.status, RollbackStatus.FAILED)
        self.assertIn("异常", result.detail)

    def test_post_rollback_health_failure(self):
        # restore 成功但健康检查失败（注入失败的健康检查）
        def _restore_then_health_fail(snapshot, ctx):
            return True, "回滚动作成功"
        mgr = RollbackManager(
            restore_action=_restore_then_health_fail, post_rollback_health=True)
        # 覆盖 _health_check 为失败
        mgr._health_check = lambda target: (False, "/health 仍不可用")
        mgr.snapshot(_make_plan("v1.0.0"))
        result = mgr.rollback_to_latest()
        self.assertEqual(result.status, RollbackStatus.FAILED)
        self.assertIn("健康检查", result.detail)
        # 压回栈顶
        self.assertEqual(len(mgr.snapshots), 1)

    def test_n_less_than_one_skipped(self):
        mgr = RollbackManager()
        result = mgr.rollback(n=0)
        self.assertEqual(result.status, RollbackStatus.SKIPPED)


class TestResultExport(unittest.TestCase):
    """结果序列化。"""

    def test_result_to_dict(self):
        mgr = RollbackManager(restore_action=_restore(True))
        mgr.snapshot(_make_plan("v1.0.0"))
        result = mgr.rollback_to_latest()
        d = result.to_dict()
        self.assertEqual(d["status"], "success")
        self.assertTrue(d["succeeded"])
        self.assertEqual(d["target"]["version"], "v1.0.0")

    def test_point_to_dict(self):
        p = RollbackPoint(release="r", namespace="ns", version="v1",
                          chart="c", values_hash="abc123", reason="x")
        d = p.to_dict()
        self.assertEqual(d["version"], "v1")
        self.assertEqual(d["values_hash"], "abc123")


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