prole/tests/installer/test_milestones.py

import unittest
from unittest.mock import MagicMock, patch, ANY
from pathlib import Path
from installer.core.milestones import (
    DependenciesMilestone, NetworkScanMilestone, EnvSetupMilestone, 
    SecretManagementMilestone, ClusterLifecycleMilestone,
    InitializationScriptsMilestone, DeploymentMilestone
)
from installer.state import InstallerState
from installer.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('installer.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('installer.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('installer.config.get_resource_path')
    @patch('subprocess.Popen')
    def test_network_scan_milestone(self, mock_popen, 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
        
        mock_process = MagicMock()
        mock_process.returncode = 0
        mock_process.stdout = MagicMock()
        mock_process.stdout.readline.side_effect = ["Scanning...", "KDC is: 10.0.0.1", ""]
        mock_process.poll.side_effect = [None, None, 0]
        mock_process.wait.return_value = 0
        mock_popen.return_value = mock_process
        
        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('installer.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_prole_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('installer.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()