Unblu MCP Server

"""Tests for the Kubernetes connection provider."""

import socket
import subprocess
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from pytest_subprocess import FakeProcess

from unblu_mcp._internal.exceptions import ConfigurationError
from unblu_mcp._internal.providers_k8s import (
    K8sConnectionProvider,
    K8sEnvironmentConfig,
    _get_default_environments,
    detect_environment_from_context,
)

# Test environments for use in tests
TEST_ENVIRONMENTS = {
    "dev": K8sEnvironmentConfig(name="dev", local_port=8084, namespace="unblu-dev"),
    "test1": K8sEnvironmentConfig(name="test1", local_port=8085, namespace="unblu-test1"),
    "test2": K8sEnvironmentConfig(name="test2", local_port=8087, namespace="unblu-test2"),
    "prod": K8sEnvironmentConfig(name="prod", local_port=8086, namespace="unblu-prod"),
}


class TestK8sEnvironmentConfig:
    """Tests for K8sEnvironmentConfig dataclass."""

    def test_default_values(self) -> None:
        """Config has sensible defaults."""
        config = K8sEnvironmentConfig(name="test", local_port=8080, namespace="test-ns")
        assert config.service == "haproxy"
        assert config.service_port == 8080
        assert config.api_path == "/app/rest/v4"

    def test_custom_values(self) -> None:
        """Config accepts custom values."""
        config = K8sEnvironmentConfig(
            name="custom",
            local_port=9000,
            namespace="custom-ns",
            service="nginx",
            service_port=443,
            api_path="/api/v1",
        )
        assert config.name == "custom"
        assert config.local_port == 9000
        assert config.namespace == "custom-ns"
        assert config.service == "nginx"
        assert config.service_port == 443
        assert config.api_path == "/api/v1"


class TestGetDefaultEnvironments:
    """Tests for environment loading from config files."""

    def test_loads_environments_from_example_config(self) -> None:
        """Loads environments from example config when user config doesn't exist."""
        environments = _get_default_environments()
        # Should load from k8s_environments.example.yaml
        assert len(environments) > 0

    def test_environments_have_unique_ports(self) -> None:
        """Each environment has a unique local port."""
        environments = _get_default_environments()
        if environments:  # Only test if environments are loaded
            ports = [env.local_port for env in environments.values()]
            assert len(ports) == len(set(ports))

    def test_environment_configs_are_valid(self) -> None:
        """All loaded environments have required fields."""
        environments = _get_default_environments()
        for name, config in environments.items():
            assert config.name == name
            assert config.local_port > 0
            assert config.namespace


class TestK8sConnectionProvider:
    """Tests for K8sConnectionProvider."""

    def test_init_with_string_environment(self) -> None:
        """Provider initializes with string environment name."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        assert provider.environment == "dev"
        assert provider.local_port == 8084

    def test_init_with_config_object(self) -> None:
        """Provider initializes with K8sEnvironmentConfig object."""
        config = K8sEnvironmentConfig(name="custom", local_port=9999, namespace="custom-ns")
        provider = K8sConnectionProvider(environment=config)
        assert provider.environment == "custom"
        assert provider.local_port == 9999

    def test_init_unknown_environment_raises(self) -> None:
        """Provider raises ValueError for unknown environment."""
        with pytest.raises(ValueError, match="Unknown environment 'invalid'"):
            K8sConnectionProvider(environment="invalid", environments=TEST_ENVIRONMENTS)

    def test_init_custom_environments(self) -> None:
        """Provider accepts custom environment configurations."""
        custom_envs = {
            "dev": K8sEnvironmentConfig(name="dev", local_port=8000, namespace="dev-ns"),
        }
        provider = K8sConnectionProvider(environment="dev", environments=custom_envs)
        assert provider.environment == "dev"
        assert provider.local_port == 8000

    def test_get_config_returns_connection_config(self) -> None:
        """get_config returns properly configured ConnectionConfig."""
        provider = K8sConnectionProvider(
            environment="dev",
            trusted_user_id="testuser",
            trusted_user_role="ADMIN",
            environments=TEST_ENVIRONMENTS,
        )
        config = provider.get_config()

        assert config.base_url == "http://localhost:8084/app/rest/v4"
        assert config.headers["x-unblu-trusted-user-id"] == "testuser"
        assert config.headers["x-unblu-trusted-user-role"] == "ADMIN"

    def test_get_config_default_trusted_headers(self) -> None:
        """get_config uses default trusted headers."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        config = provider.get_config()

        assert config.headers["x-unblu-trusted-user-id"] == "superadmin"
        assert config.headers["x-unblu-trusted-user-role"] == "SUPER_ADMIN"

    @pytest.mark.asyncio
    async def test_setup_port_already_in_use(self) -> None:
        """setup() reuses existing port if port is already in use."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        with patch.object(provider, "_is_port_in_use", return_value=True):
            await provider.setup()
            # Should not start a new process and should not own the port-forward
            assert provider._port_forward_process is None
            assert provider._owns_port_forward is False

    @pytest.mark.asyncio
    async def test_setup_kubectl_not_found(self) -> None:
        """setup() raises ConfigurationError if kubectl not found."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        with (
            patch.object(provider, "_is_port_in_use", return_value=False),
            patch("shutil.which", return_value=None),
            pytest.raises(ConfigurationError, match="kubectl not found"),
        ):
            await provider.setup()

    @pytest.mark.asyncio
    async def test_setup_kubectl_auth_failure(self, fp: FakeProcess) -> None:
        """setup() raises ConfigurationError if kubectl is not authenticated."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        # Register auth check failure
        fp.register(
            ["kubectl", "auth", "can-i", "get", "pods", "-n", "unblu-dev"],
            returncode=1,
            stderr="error: You must be logged in to the server",
        )

        with (
            patch.object(provider, "_is_port_in_use", return_value=False),
            patch("shutil.which", return_value="/usr/bin/kubectl"),
            pytest.raises(ConfigurationError, match="kubectl is not authenticated"),
        ):
            await provider.setup()

    @pytest.mark.asyncio
    async def test_setup_port_forward_fails_early(self, fp: FakeProcess) -> None:
        """setup() raises ConfigurationError with stderr if port-forward process dies early."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        # Register successful auth check
        fp.register(
            ["kubectl", "auth", "can-i", "get", "pods", "-n", "unblu-dev"],
            returncode=0,
        )
        # Register failing port-forward
        fp.register(
            ["kubectl", "port-forward", "-n", "unblu-dev", "svc/haproxy", "8084:8080"],
            returncode=1,
            stderr="error: services 'haproxy' not found",
        )

        with (
            patch.object(provider, "_is_port_in_use", return_value=False),
            patch("shutil.which", return_value="/usr/bin/kubectl"),
            patch("asyncio.sleep", new_callable=AsyncMock),
            pytest.raises(ConfigurationError, match="kubectl port-forward failed"),
        ):
            await provider.setup()

    @pytest.mark.asyncio
    async def test_setup_starts_port_forward(self, fp: FakeProcess) -> None:
        """setup() starts kubectl port-forward process when port is not in use."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        # First call returns False (port not in use), subsequent calls return True (port ready)
        port_check_results = [False, True]

        # Register successful auth check
        fp.register(
            ["kubectl", "auth", "can-i", "get", "pods", "-n", "unblu-dev"],
            returncode=0,
        )
        # Register port-forward (keeps running)
        fp.register(
            ["kubectl", "port-forward", "-n", "unblu-dev", "svc/haproxy", "8084:8080"],
            returncode=0,
        )

        with (
            patch.object(provider, "_is_port_in_use", side_effect=port_check_results),
            patch("shutil.which", return_value="/usr/bin/kubectl"),
        ):
            await provider.setup()

            assert provider._owns_port_forward is True
            # Verify the port-forward command was called
            assert fp.call_count(["kubectl", "port-forward", fp.any()]) == 1

    @pytest.mark.asyncio
    async def test_setup_timeout_kills_process(self, fp: FakeProcess) -> None:
        """setup() kills process and raises if port never becomes available."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        # Register successful auth check
        fp.register(
            ["kubectl", "auth", "can-i", "get", "pods", "-n", "unblu-dev"],
            returncode=0,
        )
        # Register port-forward that keeps running (simulated by callback)
        fp.register(
            ["kubectl", "port-forward", "-n", "unblu-dev", "svc/haproxy", "8084:8080"],
            returncode=0,
            wait=0.1,  # Small delay to simulate running process
        )

        with (
            patch.object(provider, "_is_port_in_use", return_value=False),
            patch("shutil.which", return_value="/usr/bin/kubectl"),
            patch("asyncio.sleep", new_callable=AsyncMock),
            pytest.raises(ConfigurationError, match="Port-forward timed out"),
        ):
            await provider.setup()

    @pytest.mark.asyncio
    async def test_teardown_terminates_process_when_owner(self) -> None:
        """teardown() terminates the port-forward process only when we own it."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        mock_process = MagicMock()
        provider._port_forward_process = mock_process
        provider._owns_port_forward = True

        await provider.teardown()

        mock_process.terminate.assert_called_once()
        mock_process.wait.assert_called_once_with(timeout=5)
        assert provider._port_forward_process is None

    @pytest.mark.asyncio
    async def test_teardown_kills_on_timeout(self) -> None:
        """teardown() kills process if terminate times out."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        mock_process = MagicMock()
        mock_process.wait.side_effect = subprocess.TimeoutExpired(cmd="kubectl", timeout=5)
        provider._port_forward_process = mock_process
        provider._owns_port_forward = True

        await provider.teardown()

        mock_process.terminate.assert_called_once()
        mock_process.kill.assert_called_once()
        assert provider._port_forward_process is None

    @pytest.mark.asyncio
    async def test_teardown_no_process_when_owner(self) -> None:
        """teardown() handles case where no process was started."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        provider._port_forward_process = None
        provider._owns_port_forward = True

        # Should not raise
        await provider.teardown()

    @pytest.mark.asyncio
    async def test_teardown_skips_cleanup_when_not_owner(self) -> None:
        """teardown() skips cleanup when we don't own the port-forward."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        mock_process = MagicMock()
        provider._port_forward_process = mock_process
        provider._owns_port_forward = False

        await provider.teardown()

        # Should not terminate the process since we don't own it
        mock_process.terminate.assert_not_called()
        # Process reference should remain (it's owned by another instance)
        assert provider._port_forward_process == mock_process

    @pytest.mark.asyncio
    async def test_setup_reuses_existing_port(self) -> None:
        """setup() reuses existing port-forward when port is already in use."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        with patch.object(provider, "_is_port_in_use", return_value=True):
            await provider.setup()

            assert provider._owns_port_forward is False
            assert provider._port_forward_process is None

    @pytest.mark.asyncio
    async def test_health_check_delegates_to_port_check(self) -> None:
        """health_check() returns port availability status."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        with patch.object(provider, "_is_port_in_use", return_value=True):
            assert await provider.health_check() is True

        with patch.object(provider, "_is_port_in_use", return_value=False):
            assert await provider.health_check() is False

    def test_is_port_in_use_available_port(self) -> None:
        """_is_port_in_use returns False for available port."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        # Use a high port that's unlikely to be in use
        provider._env_config = K8sEnvironmentConfig(name="test", local_port=59999, namespace="test")
        assert provider._is_port_in_use() is False

    def test_is_port_in_use_bound_port(self) -> None:
        """_is_port_in_use returns True for bound port."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        # Bind a port temporarily
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("localhost", 0))
            port = s.getsockname()[1]
            s.listen(1)

            provider._env_config = K8sEnvironmentConfig(name="test", local_port=port, namespace="test")
            assert provider._is_port_in_use() is True

    @pytest.mark.asyncio
    async def test_ensure_connection_does_nothing_when_port_available(self) -> None:
        """ensure_connection() does nothing when port is already available."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)

        with patch.object(provider, "_is_port_in_use", return_value=True):
            await provider.ensure_connection()
            # Should not have started any process
            assert provider._port_forward_process is None

    @pytest.mark.asyncio
    async def test_ensure_connection_restarts_dead_port_forward(self, fp: FakeProcess) -> None:
        """ensure_connection() restarts port-forward when our process died."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        provider._owns_port_forward = True

        # Simulate a dead process
        mock_process = MagicMock()
        mock_process.poll.return_value = 1  # Process exited with code 1
        provider._port_forward_process = mock_process

        # Port check returns False (not available), then True (after restart)
        port_check_results = [False, True]

        # Register successful auth check
        fp.register(
            ["kubectl", "auth", "can-i", "get", "pods", "-n", "unblu-dev"],
            returncode=0,
        )
        # Register port-forward
        fp.register(
            ["kubectl", "port-forward", "-n", "unblu-dev", "svc/haproxy", "8084:8080"],
            returncode=0,
        )

        with (
            patch.object(provider, "_is_port_in_use", side_effect=port_check_results),
            patch("shutil.which", return_value="/usr/bin/kubectl"),
        ):
            await provider.ensure_connection()

            # Should have started a new port-forward
            assert provider._owns_port_forward is True
            assert fp.call_count(["kubectl", "port-forward", fp.any()]) == 1

    @pytest.mark.asyncio
    async def test_ensure_connection_kills_malfunctioning_process(self, fp: FakeProcess) -> None:
        """ensure_connection() kills alive but malfunctioning port-forward."""
        provider = K8sConnectionProvider(environment="dev", environments=TEST_ENVIRONMENTS)
        provider._owns_port_forward = True

        # Simulate alive but malfunctioning process (port not available)
        mock_process = MagicMock()
        mock_process.poll.return_value = None  # Process is alive
        provider._port_forward_process = mock_process

        # Port check returns False (not available), then True (after restart)
        port_check_results = [False, True]

        # Register successful auth check
        fp.register(
            ["kubectl", "auth", "can-i", "get", "pods", "-n", "unblu-dev"],
            returncode=0,
        )
        # Register port-forward
        fp.register(
            ["kubectl", "port-forward", "-n", "unblu-dev", "svc/haproxy", "8084:8080"],
            returncode=0,
        )

        with (
            patch.object(provider, "_is_port_in_use", side_effect=port_check_results),
            patch("shutil.which", return_value="/usr/bin/kubectl"),
        ):
            await provider.ensure_connection()

            # Should have killed the malfunctioning process
            mock_process.kill.assert_called_once()
            mock_process.wait.assert_called_once_with(timeout=5)
            # Should have started a new port-forward
            assert provider._owns_port_forward is True
            assert fp.call_count(["kubectl", "port-forward", fp.any()]) == 1


class TestDetectEnvironmentFromContext:
    """Tests for detect_environment_from_context function."""

    def test_detects_dev_environment(self) -> None:
        """Detects dev from context name."""
        with (
            patch("subprocess.run") as mock_run,
            patch(
                "unblu_mcp._internal.providers_k8s._get_default_environments",
                return_value=TEST_ENVIRONMENTS,
            ),
        ):
            mock_run.return_value = MagicMock(stdout="cluster-dev-context\n")
            assert detect_environment_from_context() == "dev"

    def test_detects_environment_with_suffix(self) -> None:
        """Detects environment from context ending with env name."""
        with (
            patch("subprocess.run") as mock_run,
            patch(
                "unblu_mcp._internal.providers_k8s._get_default_environments",
                return_value=TEST_ENVIRONMENTS,
            ),
        ):
            mock_run.return_value = MagicMock(stdout="my-cluster-prod\n")
            assert detect_environment_from_context() == "prod"

    def test_detects_environment_with_dash_pattern(self) -> None:
        """Detects environment from -env- pattern in context."""
        with (
            patch("subprocess.run") as mock_run,
            patch(
                "unblu_mcp._internal.providers_k8s._get_default_environments",
                return_value=TEST_ENVIRONMENTS,
            ),
        ):
            mock_run.return_value = MagicMock(stdout="prefix-test1-suffix\n")
            assert detect_environment_from_context() == "test1"

    def test_returns_none_for_unknown_context(self) -> None:
        """Returns None if context doesn't match any environment."""
        with (
            patch("subprocess.run") as mock_run,
            patch(
                "unblu_mcp._internal.providers_k8s._get_default_environments",
                return_value=TEST_ENVIRONMENTS,
            ),
        ):
            mock_run.return_value = MagicMock(stdout="unknown-cluster\n")
            assert detect_environment_from_context() is None

    def test_returns_none_on_subprocess_error(self) -> None:
        """Returns None if kubectl command fails."""
        with patch("subprocess.run") as mock_run:
            mock_run.side_effect = subprocess.CalledProcessError(1, "kubectl")
            assert detect_environment_from_context() is None

    def test_returns_none_if_kubectl_not_found(self) -> None:
        """Returns None if kubectl is not installed."""
        with patch("subprocess.run") as mock_run:
            mock_run.side_effect = FileNotFoundError()
            assert detect_environment_from_context() is None