prole/tests/installer/test_milestones.py

import unittest
from unittest.mock import MagicMock, patch, ANY
from pathlib import Path
from knoe.core.milestones import (
    DependenciesMilestone,
    NetworkScanMilestone,
    EnvSetupMilestone,
    SecretManagementMilestone,
    ClusterLifecycleMilestone,
    InitializationScriptsMilestone,
    DeploymentMilestone,
)
from knoe.state import InstallerState
from knoe.core.controller import ProleController


class TestMilestones(unittest.TestCase):
    def setUp(self):
        self.state = InstallerState()
        self.state.inputs["dependencies.auto_install_missing"] = "True"
        self.state.inputs["network_scan.run"] = "True"
        self.state.controller = MagicMock(spec=ProleController)

    @patch("knoe.config.get_dep_info")
    @patch("subprocess.run")
    def test_dependencies_milestone_all_installed(self, mock_run, mock_get_dep_info):
        # Mock all dependencies as installed
        mock_get_dep_info.return_value = (True, "/usr/bin/foo", "1.0.0")

        milestone = DependenciesMilestone()
        progress_mock = MagicMock()

        milestone.execute(self.state, progress=progress_mock)

        self.assertEqual(
            self.state.config_data.get("Dependencies", {}).get("STATUS"),
            "All installed",
        )
        progress_mock.assert_called()

    @patch("knoe.config.get_dep_info")
    @patch("subprocess.Popen")
    def test_dependencies_milestone_install_missing(
        self, mock_popen, mock_get_dep_info
    ):
        # Mock one missing dependency
        def side_effect(dep):
            if dep["id"] == "brew":
                if mock_popen.called:
                    return (True, "/usr/local/bin/brew", "4.0.0")
                return (False, None, None)
            return (True, "/usr/bin/python", "3.10.0")

        mock_get_dep_info.side_effect = side_effect

        # Mock subprocess successful run
        mock_process = MagicMock()
        mock_process.returncode = 0
        mock_process.stdout = MagicMock()
        mock_process.stdout.readline.side_effect = ["Installing...", ""]
        mock_process.poll.side_effect = [None, 0]
        mock_process.wait.return_value = 0
        mock_popen.return_value = mock_process

        # Ensure brew install is enabled in state
        self.state.inputs["dependencies.brew.install"] = "True"

        milestone = DependenciesMilestone()
        milestone.execute(self.state)

        self.assertEqual(
            self.state.config_data.get("Dependencies", {}).get("STATUS"),
            "All installed",
        )
        # Check that it tried to install brew
        mock_popen.assert_called()

    @patch("knoe.config.get_resource_path")
    @patch("knoe.core.milestones.run_streaming_cmd")
    def test_network_scan_milestone(self, mock_run_streaming_cmd, mock_get_resource):
        mock_binary = MagicMock()
        mock_binary.exists.return_value = True
        mock_binary.__str__.return_value = "/tmp/prole-agent"
        mock_get_resource.return_value = mock_binary

        def _stream_output(cmd, cwd=None, on_stdout=None):
            if on_stdout:
                on_stdout("Scanning...\n")
                on_stdout("KDC is: 10.0.0.1\n")
            return 0

        mock_run_streaming_cmd.side_effect = _stream_output

        milestone = NetworkScanMilestone()
        milestone.execute(self.state)

        self.assertEqual(self.state.inputs.get("kerberos_config.kdc"), "10.0.0.1")
        self.assertEqual(
            self.state.config_data.get("Network", {}).get("KDC_AUTO_DETECTED"),
            "10.0.0.1",
        )

    def test_env_setup_milestone(self):
        self.state.inputs["env_setup.PROLE_HOME"] = "/custom/home"
        self.state.inputs["env_setup.NAMESPACE"] = "test-ns"

        milestone = EnvSetupMilestone()
        milestone.execute(self.state)

        self.assertEqual(
            self.state.config_data.get("System Environment", {}).get("PROLE_HOME"),
            "/custom/home",
        )
        self.assertEqual(
            self.state.config_data.get("System Environment", {}).get("NAMESPACE"),
            "test-ns",
        )

    @patch("knoe.core.milestones.inst_config._resolve_secret_value")
    def test_secret_management_milestone(self, mock_resolve):
        mock_resolve.side_effect = lambda x: f"resolved-{x}" if "secret" in x else x

        self.state.inputs["init_password.db_password"] = "secret-pw"

        milestone = SecretManagementMilestone()
        milestone.execute(self.state)

        self.assertEqual(
            self.state.inputs.get("init_password.db_password"), "resolved-secret-pw"
        )
        self.assertEqual(
            self.state.inputs.get("init_password.db_password_confirm"),
            "resolved-secret-pw",
        )

    @patch("subprocess.run")
    def test_cluster_lifecycle_milestone_dev(self, mock_subprocess_run):
        self.state.inputs["init_cluster.cluster_env"] = "dev"
        self.state.inputs["init_cluster.start_cluster"] = "True"
        self.state.controller.run_script.return_value = 0
        self.state.controller.check_docker_running.return_value = True
        self.state.controller.get_knoe_db_version.return_value = "1.0.0"

        # Generic success result for k3d kubeconfig merge / misc calls
        mock_ok = MagicMock()
        mock_ok.returncode = 0
        mock_ok.stdout = ""

        # Mock k3d cluster list showing cluster exists and running
        mock_list_result = MagicMock()
        mock_list_result.stdout = "prole-dev-cluster  1/1  0/2  running"
        mock_list_result.returncode = 0
        # Mock docker image inspect (image not found locally)
        mock_inspect_result = MagicMock()
        mock_inspect_result.returncode = 1
        # Mock kubectl config use-context
        mock_ctx_result = MagicMock()
        mock_ctx_result.returncode = 0

        # Use a default return so extra calls (k3d kubeconfig merge, etc.) don't exhaust side_effect
        mock_subprocess_run.return_value = mock_ok
        mock_subprocess_run.side_effect = None

        def _side_effect(*args, **kwargs):
            cmd = args[0] if args else kwargs.get("args", [])
            if cmd and cmd[0] == "k3d":
                if "cluster" in cmd and "list" in cmd:
                    return mock_list_result
                # k3d kubeconfig merge, k3d cluster start, k3d image import, etc.
                return mock_ok
            if cmd and cmd[0] == "kubectl":
                return mock_ctx_result
            if cmd and cmd[0] == "docker":
                return mock_inspect_result
            return mock_ok

        mock_subprocess_run.side_effect = _side_effect

        milestone = ClusterLifecycleMilestone()
        milestone.execute(self.state)

        # Verify k3d cluster list was called
        mock_subprocess_run.assert_any_call(
            ["k3d", "cluster", "list", "--no-headers"], capture_output=True, text=True
        )

    @patch("knoe.core.milestones.InitializationScriptsMilestone._run_cmd")
    @patch("subprocess.run")
    def test_init_scripts_milestone(self, mock_subprocess_run, mock_run_cmd):
        self.state.controller.run_script.return_value = 0
        # Mock _run_cmd to return 0 (secrets exist)
        mock_run_cmd.return_value = 0

        milestone = InitializationScriptsMilestone()
        milestone.execute(self.state)

        # Verify current script list is called
        self.state.controller.run_script.assert_any_call(
            "init_common_services.sh", args=ANY, env=ANY, on_line=ANY
        )
        self.state.controller.run_script.assert_any_call(
            "init_cloudnative_pg.sh", args=ANY, env=ANY, on_line=ANY
        )
        self.state.controller.run_script.assert_any_call(
            "init_monitoring.sh", args=ANY, env=ANY, on_line=ANY
        )

    def test_deployment_milestone(self):
        self.state.controller.run_script.return_value = 0

        milestone = DeploymentMilestone()
        milestone.execute(self.state)

        self.state.controller.run_script.assert_any_call(
            "init_cloudnative_pg.sh", args=ANY, env=ANY, on_line=ANY
        )


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