prole/tests/installer/test_milestones.py

import unittest
from unittest.mock import MagicMock, patch
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')

    def test_cluster_lifecycle_milestone_dev(self):
        self.state.inputs['init_cluster.cluster_env'] = 'dev'
        self.state.inputs['init_cluster.start_cluster'] = 'True'
        self.state.controller.run_script.return_value = 0
        
        milestone = ClusterLifecycleMilestone()
        milestone.execute(self.state)
        
        from unittest.mock import ANY
        self.state.controller.run_script.assert_any_call("init_k3d.sh", args=["start"], env=ANY)

    def test_init_scripts_milestone(self):
        self.state.controller.run_script.return_value = 0
        
        milestone = InitializationScriptsMilestone()
        milestone.execute(self.state)
        
        from unittest.mock import ANY
        self.state.controller.run_script.assert_any_call("init_authority.sh", args=["update"], env=ANY)
        self.state.controller.run_script.assert_any_call("init_opentofu.sh", args=["-n", "default", "update"], env=ANY)

    def test_deployment_milestone(self):
        self.state.controller.run_script.return_value = 0
        
        milestone = DeploymentMilestone()
        milestone.execute(self.state)
        
        from unittest.mock import ANY
        self.state.controller.run_script.assert_any_call("deploy_pipeline.sh", args=["--mode", "local", "apply"], env=ANY)

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