agent_deply/tests/test_agent_flow.py

from pathlib import Path

from pam_deploy_graph.agent import PamDeployAgent
from pam_deploy_graph.checkpoint_store import load_agent_state
from pam_deploy_graph.constants import GLOBAL_ACTION_SEQUENCE
from pam_deploy_graph.fake_runner import FakeActionRunner


PARAMS = {
    "HOME_BASE_URL": "https://pam.home.example.com",
    "CLIENT_ID": "client",
    "CLIENT_SECRET": "secret",
    "AIRPORT_CODE": "HET",
    "APP_NAME": "PAM",
    "MODULE_NAME": "Node",
    "VERSION_NUMBER": "2.0.5",
    "ZIP_FILE_PATH": "C:/pkg.zip",
}


def test_run_deploy_flow_success(tmp_path: Path):
    agent = PamDeployAgent(fake_runner=FakeActionRunner())
    state = agent.create_state(
        params=PARAMS,
        execution_strategy="fake",
        config_path=str(tmp_path / "config.txt"),
    )

    agent.run_deploy_flow(state)

    assert state.pending_confirmation == ""
    assert set(state.ip_states) == {"192.168.1.10", "192.168.1.11"}
    assert all(item["status"] == "SUCCESS" for item in state.ip_states.values())


def test_run_deploy_flow_stops_on_verify_failure(tmp_path: Path):
    fake = FakeActionRunner(
        {
            "verify-ip:192.168.1.10": {
                "ACTION": "verify-ip",
                "IP": "192.168.1.10",
                "SUCCESS": "false",
                "MESSAGE": "health check failed",
            }
        }
    )
    agent = PamDeployAgent(fake_runner=fake)
    state = agent.create_state(
        params=PARAMS,
        execution_strategy="fake",
        config_path=str(tmp_path / "config.txt"),
    )

    agent.run_deploy_flow(state)

    assert state.pending_confirmation == "rollback-ip:192.168.1.10"
    assert state.ip_states["192.168.1.10"]["status"] == "FAILED"
    assert state.ip_states["192.168.1.10"]["rollback_status"] == "PENDING_AGENT_CONFIRMATION"
    assert "192.168.1.11" not in state.ip_states
    assert any(event["type"] == "CONFIRMATION_REQUIRED" for event in state.events)


def test_confirm_pending_rollback_runs_rollback_and_resume_continues(tmp_path: Path):
    fake = FakeActionRunner(
        {
            "verify-ip:192.168.1.10": {
                "ACTION": "verify-ip",
                "IP": "192.168.1.10",
                "SUCCESS": "false",
                "MESSAGE": "health check failed",
            }
        }
    )
    agent = PamDeployAgent(fake_runner=fake)
    state = agent.create_state(
        params=PARAMS,
        execution_strategy="fake",
        config_path=str(tmp_path / "config.txt"),
    )

    agent.run_deploy_flow(state)
    request = agent.build_confirmation_request(state)
    agent.confirm_pending(state, approved=True)
    agent.run_deploy_flow(state)

    assert request["type"] == "rollback-ip"
    assert state.pending_confirmation == ""
    assert state.ip_states["192.168.1.10"]["rollback_status"] == "ROLLBACK_DONE"
    assert state.ip_states["192.168.1.11"]["status"] == "SUCCESS"
    assert any(call[0] == "rollback-ip" for call in fake.calls)


def test_failed_rollback_keeps_confirmation_pending(tmp_path: Path):
    fake = FakeActionRunner(
        {
            "verify-ip:192.168.1.10": {
                "ACTION": "verify-ip",
                "IP": "192.168.1.10",
                "SUCCESS": "false",
                "MESSAGE": "health check failed",
            },
            "rollback-ip:192.168.1.10": {
                "_fail": True,
                "ACTION": "rollback-ip",
                "IP": "192.168.1.10",
                "MESSAGE": "rollback failed",
            },
        }
    )
    agent = PamDeployAgent(fake_runner=fake)
    state = agent.create_state(
        params=PARAMS,
        execution_strategy="fake",
        config_path=str(tmp_path / "config.txt"),
    )

    agent.run_deploy_flow(state)
    agent.confirm_pending(state, approved=True)

    assert state.pending_confirmation == "rollback-ip:192.168.1.10"
    assert state.ip_states["192.168.1.10"]["rollback_status"] == "ROLLBACK_FAILED"


def test_checkpoint_resume_skips_completed_global_and_success_ip(tmp_path: Path):
    checkpoint = tmp_path / "checkpoint.json"
    fake = FakeActionRunner()
    agent = PamDeployAgent(fake_runner=fake)
    state = agent.create_state(
        params=PARAMS,
        execution_strategy="fake",
        config_path=str(tmp_path / "config.txt"),
        checkpoint_path=str(checkpoint),
    )
    state.completed_global_steps = list(GLOBAL_ACTION_SEQUENCE)
    state.online_ips = ["192.168.1.10", "192.168.1.11"]
    state.target_ips = ["192.168.1.10", "192.168.1.11"]
    state.ip_states["192.168.1.10"] = {
        "status": "SUCCESS",
        "completed_steps": ["upgrade-ip", "poll-upgrade-progress", "start-ip", "verify-ip", "download-log"],
        "failed_stage": "",
        "failure_reason": "",
        "rollback_status": "ROLLBACK_NOT_RUN",
        "rollback_stop_first": False,
        "log_file": "logs/fake.zip",
    }

    agent.run_deploy_flow(state)
    loaded = load_agent_state(checkpoint)

    called_actions = [call[0] for call in fake.calls]
    assert "get-token" not in called_actions
    assert all(call[1].get("ip") != "192.168.1.10" for call in fake.calls)
    assert loaded.ip_states["192.168.1.11"]["status"] == "SUCCESS"