import pytest
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import sys

class TestTransport:
    """Test the transport functionality."""
    
    @pytest.mark.asyncio
    async def test_transport_selection(self):
        """Test that the server selects the correct transport based on the argument."""
        # Import the main function here to avoid issues with the mocking
        from mcp_server_hello_world.server import main
        
        # Mock the run_stdio_async method
        with patch('mcp.server.fastmcp.FastMCP.run_stdio_async') as mock_run_stdio:
            
            # Set up the mocks
            mock_run_stdio.return_value = AsyncMock()
            
            # Start the server with stdio transport
            task = asyncio.create_task(main("stdio", "0.0.0.0", 8080))
            
            # Wait for the server to start
            await asyncio.sleep(0.1)
            
            # Check that run_stdio_async was called
            mock_run_stdio.assert_called_once()
            
            # Cancel the task
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
    
    @pytest.mark.asyncio
    async def test_remote_transport(self):
        """Test that the server can be started with remote transport."""
        # Import the main function here to avoid issues with the mocking
        from mcp_server_hello_world.server import main
        
        # Mock the run_sse_async method
        with patch('mcp.server.fastmcp.FastMCP.run_sse_async') as mock_run_sse:
            
            # Set up the mocks
            mock_run_sse.return_value = AsyncMock()
            
            # Start the server with remote transport
            task = asyncio.create_task(main("remote", "0.0.0.0", 8080))
            
            # Wait for the server to start
            await asyncio.sleep(0.1)
            
            # Check that run_sse_async was called without parameters
            mock_run_sse.assert_called_once_with()
            
            # Cancel the task
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass