litellm/tests/test_litellm/proxy/client/cli/test_keys_commands.py

import json
import os
import sys
from unittest.mock import patch

import requests

sys.path.insert(
    0, os.path.abspath("../../..")
)  # Adds the parent directory to the system path



import pytest
from click.testing import CliRunner

from litellm.proxy.client.cli import cli


@pytest.fixture
def cli_runner():
    return CliRunner()


@pytest.fixture(autouse=True)
def mock_env():
    with patch.dict(
        os.environ,
        {
            "LITELLM_PROXY_URL": "http://localhost:4000",
            "LITELLM_PROXY_API_KEY": "sk-test",
        },
    ):
        yield


@pytest.fixture
def mock_keys_client():
    with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient:
        yield MockClient


def test_async_keys_list_json_format(mock_keys_client, cli_runner):
    mock_keys_client.return_value.list.return_value = {
        "keys": [
            {
                "token": "abc123",
                "key_alias": "alias1",
                "user_id": "u1",
                "team_id": "t1",
                "spend": 10.0,
            }
        ]
    }
    result = cli_runner.invoke(cli, ["keys", "list", "--format", "json"])
    assert result.exit_code == 0
    output_data = json.loads(result.output)
    assert output_data == mock_keys_client.return_value.list.return_value
    mock_keys_client.assert_called_once_with("http://localhost:4000", "sk-test")
    mock_keys_client.return_value.list.assert_called_once()


def test_async_keys_list_table_format(mock_keys_client, cli_runner):
    mock_keys_client.return_value.list.return_value = {
        "keys": [
            {
                "token": "abc123",
                "key_alias": "alias1",
                "user_id": "u1",
                "team_id": "t1",
                "spend": 10.0,
            }
        ]
    }
    result = cli_runner.invoke(cli, ["keys", "list"])
    assert result.exit_code == 0
    assert "abc123" in result.output
    assert "alias1" in result.output
    assert "u1" in result.output
    assert "t1" in result.output
    assert "10.0" in result.output
    mock_keys_client.assert_called_once_with("http://localhost:4000", "sk-test")
    mock_keys_client.return_value.list.assert_called_once()


def test_async_keys_generate_success(mock_keys_client, cli_runner):
    mock_keys_client.return_value.generate.return_value = {
        "key": "new-key",
        "spend": 100.0,
    }
    result = cli_runner.invoke(cli, ["keys", "generate", "--models", "gpt-4", "--spend", "100"])
    assert result.exit_code == 0
    assert "new-key" in result.output
    mock_keys_client.return_value.generate.assert_called_once()


def test_async_keys_delete_success(mock_keys_client, cli_runner):
    mock_keys_client.return_value.delete.return_value = {
        "status": "success",
        "deleted_keys": ["abc123"],
    }
    result = cli_runner.invoke(cli, ["keys", "delete", "--keys", "abc123"])
    assert result.exit_code == 0
    assert "success" in result.output
    assert "abc123" in result.output
    mock_keys_client.return_value.delete.assert_called_once()


def test_async_keys_list_error_handling(mock_keys_client, cli_runner):
    mock_keys_client.return_value.list.side_effect = Exception("API Error")
    result = cli_runner.invoke(cli, ["keys", "list"])
    assert result.exit_code != 0
    assert "API Error" in str(result.exception)


def test_async_keys_generate_error_handling(mock_keys_client, cli_runner):
    mock_keys_client.return_value.generate.side_effect = Exception("API Error")
    result = cli_runner.invoke(cli, ["keys", "generate", "--models", "gpt-4"])
    assert result.exit_code != 0
    assert "API Error" in str(result.exception)


def test_async_keys_delete_error_handling(mock_keys_client, cli_runner):
    import requests

    # Mock a connection error that would normally happen in CI
    mock_keys_client.return_value.delete.side_effect = requests.exceptions.ConnectionError(
        "Connection error"
    )
    result = cli_runner.invoke(cli, ["keys", "delete", "--keys", "abc123"])
    assert result.exit_code != 0
    # Check that the exception is properly propagated
    assert result.exception is not None
    # The ConnectionError should propagate since it's not caught by HTTPError handler
    # Check for connection-related keywords that appear in both mocked and real errors
    error_str = str(result.exception).lower()
    assert any(keyword in error_str for keyword in ["connection", "connect", "refused", "error"])


def test_async_keys_delete_http_error_handling(mock_keys_client, cli_runner):
    from unittest.mock import Mock

    import requests

    # Create a mock response object for HTTPError
    mock_response = Mock()
    mock_response.status_code = 400
    mock_response.json.return_value = {"error": "Bad request"}
    
    # Mock an HTTPError which should be caught by the delete command
    http_error = requests.exceptions.HTTPError("HTTP Error")
    http_error.response = mock_response
    mock_keys_client.return_value.delete.side_effect = http_error
    
    result = cli_runner.invoke(cli, ["keys", "delete", "--keys", "abc123"])
    assert result.exit_code != 0
    # HTTPError should be caught and converted to click.Abort
    assert isinstance(result.exception, SystemExit)  # click.Abort raises SystemExit


# Tests for keys import command
def test_keys_import_dry_run_success(mock_keys_client, cli_runner):
    """Test successful dry-run import showing table of keys that would be imported"""
    # Mock source client response (paginated)
    mock_source_instance = mock_keys_client.return_value
    mock_source_instance.list.side_effect = [
        {
            "keys": [
                {
                    "key_alias": "test-key-1",
                    "user_id": "user1@example.com",
                    "created_at": "2024-01-15T10:30:00Z",
                    "models": ["gpt-4"],
                    "spend": 10.0,
                },
                {
                    "key_alias": "test-key-2", 
                    "user_id": "user2@example.com",
                    "created_at": "2024-01-16T11:45:00Z",
                    "models": [],
                    "spend": 5.0,
                }
            ]
        },
        {"keys": []}  # Empty second page
    ]
    
    result = cli_runner.invoke(cli, [
        "keys", "import",
        "--source-base-url", "https://source.example.com",
        "--source-api-key", "sk-source-123",
        "--dry-run"
    ])
    
    assert result.exit_code == 0
    assert "Found 2 keys in source instance" in result.output
    assert "DRY RUN MODE" in result.output
    assert "test-key-1" in result.output
    assert "user1@example.com" in result.output
    assert "test-key-2" in result.output
    assert "user2@example.com" in result.output
    
    # Verify source client was called (pagination stops early when fewer keys than page_size)
    assert mock_source_instance.list.call_count >= 1
    mock_source_instance.list.assert_any_call(return_full_object=True, page=1, size=100)


def test_keys_import_actual_import_success(mock_keys_client, cli_runner):
    """Test successful actual import of keys"""
    # Create separate mock instances for source and destination
    with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient:
        mock_source_instance = MockClient.return_value
        mock_dest_instance = MockClient.return_value
        
        # Configure source client
        mock_source_instance.list.side_effect = [
            {
                "keys": [
                    {
                        "key_alias": "import-key-1",
                        "user_id": "user1@example.com",
                        "models": ["gpt-4"],
                        "spend": 100.0,
                        "team_id": "team-1"
                    }
                ]
            },
            {"keys": []}  # Empty second page
        ]
        
        # Configure destination client
        mock_dest_instance.generate.return_value = {
            "key": "sk-new-generated-key",
            "status": "success"
        }
        
        result = cli_runner.invoke(cli, [
            "keys", "import",
            "--source-base-url", "https://source.example.com",
            "--source-api-key", "sk-source-123"
        ])
        
        assert result.exit_code == 0
        assert "Found 1 keys in source instance" in result.output
        assert "✓ Imported key: import-key-1" in result.output
        assert "Successfully imported: 1" in result.output
        assert "Failed to import: 0" in result.output
        
        # Verify generate was called with correct parameters
        mock_dest_instance.generate.assert_called_once_with(
            models=["gpt-4"],
            spend=100.0,
            key_alias="import-key-1",
            team_id="team-1",
            user_id="user1@example.com"
        )


def test_keys_import_pagination_handling(mock_keys_client, cli_runner):
    """Test that import correctly handles pagination to get all keys"""
    mock_source_instance = mock_keys_client.return_value
    mock_source_instance.list.side_effect = [
        {"keys": [{"key_alias": f"key-{i}", "user_id": f"user{i}@example.com"} for i in range(100)]},  # Page 1: 100 keys
        {"keys": [{"key_alias": f"key-{i}", "user_id": f"user{i}@example.com"} for i in range(100, 150)]},  # Page 2: 50 keys
        {"keys": []}  # Page 3: Empty
    ]
    
    result = cli_runner.invoke(cli, [
        "keys", "import", 
        "--source-base-url", "https://source.example.com",
        "--dry-run"
    ])
    
    assert result.exit_code == 0
    assert "Fetched page 1: 100 keys" in result.output
    assert "Fetched page 2: 50 keys" in result.output
    assert "Found 150 keys in source instance" in result.output
    
    # Verify pagination calls (stops early when fewer keys than page_size)
    assert mock_source_instance.list.call_count >= 2
    mock_source_instance.list.assert_any_call(return_full_object=True, page=1, size=100)
    mock_source_instance.list.assert_any_call(return_full_object=True, page=2, size=100)


def test_keys_import_created_since_filter(mock_keys_client, cli_runner):
    """Test that --created-since filter works correctly"""
    mock_source_instance = mock_keys_client.return_value
    mock_source_instance.list.side_effect = [
        {
            "keys": [
                {
                    "key_alias": "old-key",
                    "user_id": "user1@example.com", 
                    "created_at": "2024-01-01T10:00:00Z"  # Before filter
                },
                {
                    "key_alias": "new-key",
                    "user_id": "user2@example.com",
                    "created_at": "2024-07-08T10:00:00Z"  # After filter
                }
            ]
        },
        {"keys": []}
    ]
    
    result = cli_runner.invoke(cli, [
        "keys", "import",
        "--source-base-url", "https://source.example.com", 
        "--created-since", "2024-07-07_18:19",
        "--dry-run"
    ])
    
    assert result.exit_code == 0
    assert "Filtered 2 keys to 1 keys created since 2024-07-07_18:19" in result.output
    assert "Found 1 keys in source instance" in result.output
    assert "new-key" in result.output
    assert "old-key" not in result.output


def test_keys_import_created_since_date_only_format(mock_keys_client, cli_runner):
    """Test --created-since with date-only format (YYYY-MM-DD)"""
    mock_source_instance = mock_keys_client.return_value
    mock_source_instance.list.side_effect = [
        {
            "keys": [
                {
                    "key_alias": "test-key",
                    "user_id": "user@example.com",
                    "created_at": "2024-07-08T10:00:00Z"
                }
            ]
        },
        {"keys": []}
    ]
    
    result = cli_runner.invoke(cli, [
        "keys", "import",
        "--source-base-url", "https://source.example.com",
        "--created-since", "2024-07-07",  # Date only format
        "--dry-run"
    ])
    
    assert result.exit_code == 0
    assert "Filtered 1 keys to 1 keys created since 2024-07-07" in result.output


def test_keys_import_no_keys_found(mock_keys_client, cli_runner):
    """Test handling when no keys are found in source instance"""
    mock_source_instance = mock_keys_client.return_value
    mock_source_instance.list.return_value = {"keys": []}
    
    result = cli_runner.invoke(cli, [
        "keys", "import",
        "--source-base-url", "https://source.example.com",
        "--dry-run"
    ])
    
    assert result.exit_code == 0
    assert "No keys found in source instance" in result.output


def test_keys_import_invalid_date_format(cli_runner):
    """Test error handling for invalid --created-since date format"""
    result = cli_runner.invoke(cli, [
        "keys", "import",
        "--source-base-url", "https://source.example.com",
        "--created-since", "invalid-date",
        "--dry-run"
    ])
    
    assert result.exit_code != 0
    assert "Invalid date format" in result.output
    assert "Use YYYY-MM-DD_HH:MM or YYYY-MM-DD" in result.output


def test_keys_import_source_api_error(mock_keys_client, cli_runner):
    """Test error handling when source API returns an error"""
    mock_source_instance = mock_keys_client.return_value
    mock_source_instance.list.side_effect = Exception("Source API Error")
    
    result = cli_runner.invoke(cli, [
        "keys", "import",
        "--source-base-url", "https://source.example.com",
        "--dry-run"
    ])
    
    assert result.exit_code != 0
    assert "Source API Error" in result.output


def test_keys_import_partial_failure(mock_keys_client, cli_runner):
    """Test handling when some keys fail to import"""
    with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient:
        mock_source_instance = MockClient.return_value
        mock_dest_instance = MockClient.return_value
        
        # Source returns 2 keys
        mock_source_instance.list.side_effect = [
            {
                "keys": [
                    {"key_alias": "success-key", "user_id": "user1@example.com"},
                    {"key_alias": "fail-key", "user_id": "user2@example.com"}
                ]
            },
            {"keys": []}
        ]
        
        # Destination: first succeeds, second fails
        mock_dest_instance.generate.side_effect = [
            {"key": "sk-new-key", "status": "success"},
            Exception("Import failed for this key")
        ]
        
        result = cli_runner.invoke(cli, [
            "keys", "import",
            "--source-base-url", "https://source.example.com"
        ])
        
        assert result.exit_code == 0  # Command completes even with partial failures
        assert "✓ Imported key: success-key" in result.output
        assert "✗ Failed to import key fail-key" in result.output
        assert "Successfully imported: 1" in result.output
        assert "Failed to import: 1" in result.output
        assert "Total keys processed: 2" in result.output


def test_keys_import_missing_required_source_url(cli_runner):
    """Test error when required --source-base-url is missing"""
    result = cli_runner.invoke(cli, [
        "keys", "import",
        "--dry-run"
    ])
    
    assert result.exit_code != 0
    assert "Missing option" in result.output or "required" in result.output.lower()


def test_keys_import_with_all_key_properties(mock_keys_client, cli_runner):
    """Test import preserves all key properties (models, aliases, config, etc.)"""
    with patch("litellm.proxy.client.cli.commands.keys.KeysManagementClient") as MockClient:
        mock_source_instance = MockClient.return_value
        mock_dest_instance = MockClient.return_value
        
        mock_source_instance.list.side_effect = [
            {
                "keys": [
                    {
                        "key_alias": "full-key",
                        "user_id": "user@example.com",
                        "team_id": "team-123",
                        "budget_id": "budget-456", 
                        "models": ["gpt-4", "gpt-3.5-turbo"],
                        "aliases": {"custom-model": "gpt-4"},
                        "spend": 50.0,
                        "config": {"max_tokens": 1000}
                    }
                ]
            },
            {"keys": []}
        ]
        
        mock_dest_instance.generate.return_value = {"key": "sk-imported", "status": "success"}
        
        result = cli_runner.invoke(cli, [
            "keys", "import",
            "--source-base-url", "https://source.example.com"
        ])
        
        assert result.exit_code == 0
        
        # Verify all properties were passed to generate
        mock_dest_instance.generate.assert_called_once_with(
            models=["gpt-4", "gpt-3.5-turbo"],
            aliases={"custom-model": "gpt-4"},
            spend=50.0,
            key_alias="full-key",
            team_id="team-123",
            user_id="user@example.com",
            budget_id="budget-456",
            config={"max_tokens": 1000}
        )