test_pod_commands.py

""" Test CLI pod commands """

import unittest
from unittest.mock import MagicMock, patch, mock_open

from click.testing import CliRunner
from prettytable import PrettyTable

from runpod.cli.entry import runpod_cli


class TestPodCommands(unittest.TestCase):
    """Test CLI pod commands"""

    @patch("runpod.cli.groups.pod.commands.get_pods")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    def test_list_pods(self, mock_echo, mock_get_pods):
        """
        Test list_pods
        """
        # Mock data returned by get_pods
        mock_get_pods.return_value = [
            {
                "id": "1",
                "name": "Pod1",
                "desiredStatus": "Running",
                "imageName": "Image1",
            },
            {
                "id": "2",
                "name": "Pod2",
                "desiredStatus": "Stopped",
                "imageName": "Image2",
            },
        ]

        runner = CliRunner()
        result = runner.invoke(runpod_cli, ["pod", "list"])

        # Create expected table
        assert result.exit_code == 0, result.exception
        expected_table = PrettyTable(["ID", "Name", "Status", "Image"])
        expected_table.add_row(("1", "Pod1", "Running", "Image1"))
        expected_table.add_row(("2", "Pod2", "Stopped", "Image2"))

        # Assert that click.echo was called with the correct table
        mock_echo.assert_called()

    @patch("runpod.cli.groups.pod.commands.click.prompt")
    @patch("runpod.cli.groups.pod.commands.click.confirm")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.create_pod")
    def test_create_new_pod(
        self, mock_create_pod, mock_echo, mock_confirm, mock_prompt
    ):  # pylint: disable=too-many-arguments,line-too-long
        """
        Test create_new_pod
        """
        # Mock values
        mock_confirm.return_value = True  # for the quick_launch option
        mock_prompt.return_value = "RunPod-CLI-Pod"
        mock_create_pod.return_value = {"id": "sample_id"}

        runner = CliRunner()
        result = runner.invoke(runpod_cli, ["pod", "create"])

        # Assertions
        assert result.exit_code == 0, result.exception
        mock_prompt.assert_called_once_with("Enter pod name", default="RunPod-CLI-Pod")
        mock_echo.assert_called_with("Pod sample_id has been created.")
        mock_create_pod.assert_called_with(
            "RunPod-CLI-Pod",
            "runpod/base:0.0.0",
            "NVIDIA GeForce RTX 3090",
            gpu_count=1,
            support_public_ip=True,
            ports="22/tcp",
            encrypt_volume=False,  # Accept the new default argument
        )
        mock_echo.assert_called_with("Pod sample_id has been created.")

    @patch("runpod.cli.groups.pod.commands.click.prompt")
    @patch("runpod.cli.groups.pod.commands.click.confirm")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.create_pod")
    def test_create_new_pod_with_encrypt_volume(
        self, mock_create_pod, mock_echo, mock_confirm, mock_prompt
    ):  # pylint: disable=too-many-arguments,line-too-long
        """
        Test create_new_pod with --encrypt-volume option
        """
        # Mock values
        mock_confirm.return_value = True  # for the quick_launch option
        mock_prompt.return_value = "RunPod-CLI-Pod-Encrypted"
        mock_create_pod.return_value = {"id": "encrypted_pod_id"}

        runner = CliRunner()
        result = runner.invoke(runpod_cli, ["pod", "create", "--encrypt-volume"])

        # Assertions
        assert result.exit_code == 0, result.exception
        mock_prompt.assert_called_once_with("Enter pod name", default="RunPod-CLI-Pod")
        mock_echo.assert_called_with("Pod encrypted_pod_id has been created.")
        mock_create_pod.assert_called_with(
            "RunPod-CLI-Pod-Encrypted",
            "runpod/base:0.0.0",
            "NVIDIA GeForce RTX 3090",
            gpu_count=1,
            support_public_ip=True,
            ports="22/tcp",
            encrypt_volume=True,  # This is the key assertion
        )

    @patch("runpod.cli.groups.pod.commands.click.prompt")
    @patch("runpod.cli.groups.pod.commands.click.confirm")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.create_pod")
    def test_create_new_pod_with_no_encrypt_volume(
        self, mock_create_pod, mock_echo, mock_confirm, mock_prompt
    ):  # pylint: disable=too-many-arguments,line-too-long
        """
        Test create_new_pod with --no-encrypt-volume option
        """
        # Mock values
        mock_confirm.return_value = True  # for the quick_launch option
        mock_prompt.return_value = "RunPod-CLI-Pod-No-Encryption"
        mock_create_pod.return_value = {"id": "no_encryption_pod_id"}

        runner = CliRunner()
        result = runner.invoke(runpod_cli, ["pod", "create", "--no-encrypt-volume"])

        # Assertions
        assert result.exit_code == 0, result.exception
        mock_prompt.assert_called_once_with("Enter pod name", default="RunPod-CLI-Pod")
        mock_echo.assert_called_with("Pod no_encryption_pod_id has been created.")
        mock_create_pod.assert_called_with(
            "RunPod-CLI-Pod-No-Encryption",
            "runpod/base:0.0.0",
            "NVIDIA GeForce RTX 3090",
            gpu_count=1,
            support_public_ip=True,
            ports="22/tcp",
            encrypt_volume=False,  # This is the key assertion
        )

    @patch("runpod.cli.groups.pod.commands.click.prompt")
    @patch("runpod.cli.groups.pod.commands.click.confirm")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.create_pod")
    def test_create_new_pod_default_encrypt_volume(
        self, mock_create_pod, mock_echo, mock_confirm, mock_prompt
    ):  # pylint: disable=too-many-arguments,line-too-long
        """
        Test create_new_pod with default encrypt_volume (should be False)
        """
        # Mock values
        mock_confirm.return_value = True  # for the quick_launch option
        mock_prompt.return_value = "RunPod-CLI-Pod-Default"
        mock_create_pod.return_value = {"id": "default_pod_id"}

        runner = CliRunner()
        result = runner.invoke(runpod_cli, ["pod", "create"])

        # Assertions
        assert result.exit_code == 0, result.exception
        mock_prompt.assert_called_once_with("Enter pod name", default="RunPod-CLI-Pod")
        mock_echo.assert_called_with("Pod default_pod_id has been created.")
        mock_create_pod.assert_called_with(
            "RunPod-CLI-Pod-Default",
            "runpod/base:0.0.0",
            "NVIDIA GeForce RTX 3090",
            gpu_count=1,
            support_public_ip=True,
            ports="22/tcp",
            encrypt_volume=False,  # Default should be False
        )

    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.ssh_cmd.SSHConnection")
    def test_connect_to_pod(self, mock_ssh_connection, mock_echo):
        """
        Test connect_to_pod function
        """
        pod_id = "sample_id"

        mock_ssh = MagicMock()
        mock_ssh_connection.return_value = mock_ssh

        runner = CliRunner()
        result = runner.invoke(runpod_cli, ["pod", "connect", pod_id])

        assert result.exit_code == 0, result.exception
        mock_echo.assert_called_once_with(f"Connecting to pod {pod_id}...")
        mock_ssh_connection.assert_called_once_with(pod_id)
        mock_ssh.launch_terminal.assert_called_once_with()

    @patch("runpod.cli.groups.pod.commands.os.unlink")
    @patch("runpod.cli.groups.pod.commands.tempfile.NamedTemporaryFile")
    @patch("runpod.cli.groups.pod.commands.uuid.uuid4")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.ssh_cmd.SSHConnection")
    def test_sync_pods_success(self, mock_ssh_connection, mock_echo, mock_uuid, mock_temp_file, mock_unlink):
        """
        Test sync_pods function - successful sync
        """
        # Setup mocks
        mock_uuid.return_value = MagicMock()
        mock_uuid.return_value.__str__ = MagicMock(return_value="12345678-1234-1234-1234-123456789012")
        
        mock_temp_file.return_value.__enter__.return_value.name = "/tmp/test_archive.tar.gz"
        
        # Mock SSH connections
        mock_source_ssh = MagicMock()
        mock_dest_ssh = MagicMock()
        
        # Mock SSH exec_command responses
        mock_source_ssh.ssh.exec_command.side_effect = [
            (None, MagicMock(read=lambda: b"42"), None),  # file count
            (None, MagicMock(read=lambda: b"exists"), None),  # directory exists check
            (None, MagicMock(read=lambda: b"created"), None),  # archive created check
            (None, MagicMock(read=lambda: b"1.5M"), None),  # archive size
        ]
        
        mock_dest_ssh.ssh.exec_command.return_value = (None, MagicMock(read=lambda: b"42"), None)  # dest file count
        
        # Configure SSH connection context manager
        def ssh_side_effect(pod_id):
            if pod_id == "source_pod":
                mock_source_ssh.__enter__ = MagicMock(return_value=mock_source_ssh)
                mock_source_ssh.__exit__ = MagicMock(return_value=None)
                return mock_source_ssh
            elif pod_id == "dest_pod":
                mock_dest_ssh.__enter__ = MagicMock(return_value=mock_dest_ssh)
                mock_dest_ssh.__exit__ = MagicMock(return_value=None)
                return mock_dest_ssh
            
        mock_ssh_connection.side_effect = ssh_side_effect
        
        # Mock SSH key validation
        with patch("runpod.cli.groups.ssh.functions.get_user_pub_keys") as mock_get_keys:
            mock_get_keys.return_value = [{"name": "test-key", "type": "ssh-rsa", "fingerprint": "SHA256:test"}]
            
            runner = CliRunner()
            result = runner.invoke(runpod_cli, ["pod", "sync", "source_pod", "dest_pod", "/workspace", "/workspace"])

        assert result.exit_code == 0, result.exception
        
        # Verify SSH connections were created
        assert mock_ssh_connection.call_count == 2
        mock_ssh_connection.assert_any_call("source_pod")
        mock_ssh_connection.assert_any_call("dest_pod")
        
        # Verify file operations
        mock_source_ssh.get_file.assert_called_once()
        mock_dest_ssh.put_file.assert_called_once()
        
        # Verify commands were run
        mock_source_ssh.run_commands.assert_called()
        mock_dest_ssh.run_commands.assert_called()

    @patch("runpod.cli.groups.pod.commands.uuid.uuid4")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.ssh_cmd.SSHConnection")
    def test_sync_pods_no_ssh_keys(self, mock_ssh_connection, mock_echo, mock_uuid):
        """
        Test sync_pods function - no SSH keys configured
        """
        # Setup mocks
        mock_uuid.return_value = MagicMock()
        mock_uuid.return_value.__str__ = MagicMock(return_value="12345678-1234-1234-1234-123456789012")
        
        # Mock SSH key validation - no keys found
        with patch("runpod.cli.groups.ssh.functions.get_user_pub_keys") as mock_get_keys:
            mock_get_keys.return_value = []  # No SSH keys
            
            runner = CliRunner()
            result = runner.invoke(runpod_cli, ["pod", "sync", "source_pod", "dest_pod"])

        assert result.exit_code == 0, result.exception
        
        # Verify error message was shown
        mock_echo.assert_any_call("❌ No SSH keys found in your Runpod account!")
        mock_echo.assert_any_call("🔑 To create an SSH key, run:")
        mock_echo.assert_any_call("   runpod ssh add-key")

    @patch("runpod.cli.groups.pod.commands.uuid.uuid4")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.ssh_cmd.SSHConnection")
    def test_sync_pods_source_not_found(self, mock_ssh_connection, mock_echo, mock_uuid):
        """
        Test sync_pods function - source directory not found
        """
        # Setup mocks
        mock_uuid.return_value = MagicMock()
        mock_uuid.return_value.__str__ = MagicMock(return_value="12345678-1234-1234-1234-123456789012")
        
        mock_source_ssh = MagicMock()
        mock_source_ssh.__enter__ = MagicMock(return_value=mock_source_ssh)
        mock_source_ssh.__exit__ = MagicMock(return_value=None)
        
        # Mock SSH exec_command responses - directory doesn't exist
        mock_source_ssh.ssh.exec_command.side_effect = [
            (None, MagicMock(read=lambda: b"0"), None),  # file count
            (None, MagicMock(read=lambda: b"not_found"), None),  # directory exists check
        ]
        
        mock_ssh_connection.return_value = mock_source_ssh
        
        # Mock SSH key validation
        with patch("runpod.cli.groups.ssh.functions.get_user_pub_keys") as mock_get_keys:
            mock_get_keys.return_value = [{"name": "test-key", "type": "ssh-rsa", "fingerprint": "SHA256:test"}]
            
            runner = CliRunner()
            result = runner.invoke(runpod_cli, ["pod", "sync", "source_pod", "dest_pod", "/nonexistent", "/workspace"])

        assert result.exit_code == 0, result.exception
        
        # Verify error message was shown
        mock_echo.assert_any_call("❌ Error: Source workspace /nonexistent does not exist on pod source_pod")

    @patch("runpod.cli.groups.pod.commands.uuid.uuid4")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.ssh_cmd.SSHConnection")
    def test_sync_pods_archive_creation_failed(self, mock_ssh_connection, mock_echo, mock_uuid):
        """
        Test sync_pods function - archive creation failed
        """
        # Setup mocks
        mock_uuid.return_value = MagicMock()
        mock_uuid.return_value.__str__ = MagicMock(return_value="12345678-1234-1234-1234-123456789012")
        
        mock_source_ssh = MagicMock()
        mock_source_ssh.__enter__ = MagicMock(return_value=mock_source_ssh)
        mock_source_ssh.__exit__ = MagicMock(return_value=None)
        
        # Mock SSH exec_command responses - archive creation fails
        mock_source_ssh.ssh.exec_command.side_effect = [
            (None, MagicMock(read=lambda: b"42"), None),  # file count
            (None, MagicMock(read=lambda: b"exists"), None),  # directory exists check
            (None, MagicMock(read=lambda: b"failed"), None),  # archive created check fails
        ]
        
        mock_ssh_connection.return_value = mock_source_ssh
        
        # Mock SSH key validation
        with patch("runpod.cli.groups.ssh.functions.get_user_pub_keys") as mock_get_keys:
            mock_get_keys.return_value = [{"name": "test-key", "type": "ssh-rsa", "fingerprint": "SHA256:test"}]
            
            runner = CliRunner()
            result = runner.invoke(runpod_cli, ["pod", "sync", "source_pod", "dest_pod", "/workspace", "/workspace"])

        assert result.exit_code == 0, result.exception
        
        # Verify error message was shown
        mock_echo.assert_any_call("❌ Error: Failed to create archive on source pod")

    @patch("runpod.cli.groups.pod.commands.uuid.uuid4")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.ssh_cmd.SSHConnection")
    def test_sync_pods_ssh_exception(self, mock_ssh_connection, mock_echo, mock_uuid):
        """
        Test sync_pods function - SSH connection exception
        """
        # Setup mocks
        mock_uuid.return_value = MagicMock()
        mock_uuid.return_value.__str__ = MagicMock(return_value="12345678-1234-1234-1234-123456789012")
        
        # Mock SSH connection to raise exception
        mock_ssh_connection.side_effect = Exception("SSH connection failed")
        
        # Mock SSH key validation
        with patch("runpod.cli.groups.ssh.functions.get_user_pub_keys") as mock_get_keys:
            mock_get_keys.return_value = [{"name": "test-key", "type": "ssh-rsa", "fingerprint": "SHA256:test"}]
            
            runner = CliRunner()
            result = runner.invoke(runpod_cli, ["pod", "sync", "source_pod", "dest_pod", "/workspace", "/workspace"])

        assert result.exit_code == 0, result.exception
        
        # Verify error message was shown
        mock_echo.assert_any_call("❌ Error during sync: SSH connection failed")

    @patch("runpod.cli.groups.pod.commands.uuid.uuid4")
    @patch("runpod.cli.groups.pod.commands.click.echo")
    @patch("runpod.cli.groups.pod.commands.ssh_cmd.SSHConnection")
    def test_sync_pods_default_workspace(self, mock_ssh_connection, mock_echo, mock_uuid):
        """
        Test sync_pods function - using default workspace paths
        """
        # Setup mocks
        mock_uuid.return_value = MagicMock()
        mock_uuid.return_value.__str__ = MagicMock(return_value="12345678-1234-1234-1234-123456789012")
        
        mock_source_ssh = MagicMock()
        mock_source_ssh.__enter__ = MagicMock(return_value=mock_source_ssh)
        mock_source_ssh.__exit__ = MagicMock(return_value=None)
        
        # Mock SSH exec_command responses
        mock_source_ssh.ssh.exec_command.side_effect = [
            (None, MagicMock(read=lambda: b"10"), None),  # file count
            (None, MagicMock(read=lambda: b"exists"), None),  # directory exists check
        ]
        
        mock_ssh_connection.return_value = mock_source_ssh
        
        # Mock SSH key validation
        with patch("runpod.cli.groups.ssh.functions.get_user_pub_keys") as mock_get_keys:
            mock_get_keys.return_value = [{"name": "test-key", "type": "ssh-rsa", "fingerprint": "SHA256:test"}]
            
            runner = CliRunner()
            # Test with only pod IDs (should use /workspace as default)
            result = runner.invoke(runpod_cli, ["pod", "sync", "source_pod", "dest_pod"])

        assert result.exit_code == 0, result.exception
        
        # Verify the default workspace path is used
        mock_echo.assert_any_call("🔄 Syncing from source_pod:/workspace to dest_pod:/workspace")