"""
Tests for Runpod CLI exec commands.
"""

import tempfile
import unittest
from unittest.mock import patch

import click

from runpod.cli.entry import runpod_cli


class TestExecCommands(unittest.TestCase):
    """Tests for Runpod CLI exec commands."""

    def setUp(self):
        self.runner = click.testing.CliRunner()

    def test_remote_python_with_provided_pod_id(self):
        """Tests the remote_python command when pod_id is provided directly."""
        with tempfile.NamedTemporaryFile() as temp_file, patch(
            "runpod.cli.groups.exec.commands.python_over_ssh"
        ) as mock_python_over_ssh:
            result = self.runner.invoke(
                runpod_cli,
                ["exec", "python", "--pod_id", "sample_pod_id", temp_file.name],
            )
            assert result.exit_code == 0
            mock_python_over_ssh.assert_called_with("sample_pod_id", temp_file.name)

    def test_remote_python_without_provided_pod_id_stored(self):
        """Tests the remote_python command when pod_id is retrieved from storage."""
        with tempfile.NamedTemporaryFile() as temp_file, patch(
            "runpod.cli.groups.exec.commands.python_over_ssh"
        ) as mock_python_over_ssh, patch(
            "runpod.cli.groups.exec.commands.get_session_pod",
            return_value="stored_pod_id",
        ) as mock_get_pod_id:  # pylint: disable=line-too-long
            mock_python_over_ssh.return_value = None
            result = self.runner.invoke(runpod_cli, ["exec", "python", temp_file.name])
            assert result.exit_code == 0
            mock_get_pod_id.assert_called_once()
            mock_python_over_ssh.assert_called_with("stored_pod_id", temp_file.name)

    def test_remote_python_without_provided_pod_id_prompt(self):
        """Tests the remote_python command when pod_id is prompted to user."""
        with tempfile.NamedTemporaryFile() as temp_file, patch(
            "runpod.cli.groups.exec.commands.python_over_ssh"
        ) as mock_python_over_ssh, patch(
            "runpod.cli.groups.exec.commands.get_session_pod",
            side_effect=lambda: click.prompt(
                "Please provide the pod ID", "prompted_pod_id"
            ),
        ) as mock_get_pod_id:  # pylint: disable=line-too-long
            mock_python_over_ssh.return_value = None
            result = self.runner.invoke(runpod_cli, ["exec", "python", temp_file.name])
            assert result.exit_code == 0
            mock_get_pod_id.assert_called_once()
            mock_python_over_ssh.assert_called_with("prompted_pod_id", temp_file.name)
