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

import json
import os
from unittest.mock import patch

import pytest
import requests
import responses
from click.testing import CliRunner

from litellm.proxy.client.cli import cli
from litellm.proxy.client.cli.commands.keys import keys

DESTINATION = "https://destination.example.com"
SOURCE = "https://source.example.com"


@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):

    # 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


    # 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},
        )


def _sent_requests():
    return [
        (call.request.method, call.request.url, call.request.headers.get("Authorization")) for call in responses.calls
    ]


@pytest.mark.parametrize(("api_key", "authorization"), [("sk-destination", "Bearer sk-destination"), (None, None)])
@pytest.mark.parametrize(
    ("args", "method", "path", "sent_body", "payload"),
    [
        (
            ["list", "--format", "json"],
            "GET",
            "/key/list?return_full_object=true&include_team_keys=false",
            None,
            {"keys": [{"token": "abc123", "key_alias": "alias1"}], "total_count": 1},
        ),
        (["generate", "--key-alias", "ci"], "POST", "/key/generate", {"key_alias": "ci"}, {"key": "sk-new"}),
        (
            ["delete", "--keys", "sk-old"],
            "POST",
            "/key/delete",
            {"keys": ["sk-old"], "key_aliases": None},
            {"deleted_keys": ["sk-old"]},
        ),
    ],
)
@responses.activate
def test_keys_command_calls_the_server_stored_on_the_cli_context(
    cli_runner, args, method, path, sent_body, payload, api_key, authorization
):
    responses.add(method, f"{DESTINATION}{path.partition('?')[0]}", json=payload)

    result = cli_runner.invoke(keys, args, obj={"base_url": DESTINATION, "api_key": api_key})

    assert result.exit_code == 0, result.output
    assert json.loads(result.output) == payload
    assert _sent_requests() == [(method, f"{DESTINATION}{path}", authorization)]
    assert (responses.calls[0].request.body and json.loads(responses.calls[0].request.body)) == sent_body


@pytest.mark.parametrize(("api_key", "authorization"), [("sk-destination", "Bearer sk-destination"), (None, None)])
@responses.activate
def test_keys_import_reads_the_source_and_writes_to_the_server_stored_on_the_cli_context(
    cli_runner, api_key, authorization
):
    responses.get(f"{SOURCE}/key/list", json={"keys": [{"key_alias": "imported", "user_id": "u1"}]})
    responses.post(f"{DESTINATION}/key/generate", json={"key": "sk-imported"})

    result = cli_runner.invoke(
        keys,
        ["import", "--source-base-url", SOURCE, "--source-api-key", "sk-source"],
        obj={"base_url": DESTINATION, "api_key": api_key},
    )

    assert result.exit_code == 0, result.output
    assert "Successfully imported: 1" in result.output
    assert _sent_requests() == [
        ("GET", f"{SOURCE}/key/list?page=1&size=100&return_full_object=true", "Bearer sk-source"),
        ("POST", f"{DESTINATION}/key/generate", authorization),
    ]
    assert json.loads(responses.calls[1].request.body) == {"key_alias": "imported", "user_id": "u1"}


@pytest.mark.parametrize(
    "args",
    [
        ["list"],
        ["generate", "--key-alias", "ci"],
        ["delete", "--keys", "sk-old"],
        ["import", "--source-base-url", SOURCE, "--source-api-key", "sk-source"],
    ],
)
@responses.activate
def test_keys_command_without_cli_context_values_fails_before_any_request(cli_runner, args):
    result = cli_runner.invoke(keys, args)

    assert result.exit_code == 1
    assert _sent_requests() == []