mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-17 22:48:35 +00:00
boarder7395
8e4f06583a
* Cleanup code for user cli auth, and make sure not to prompt user for team multiple times while polling * Adding tests * Cleanup normalize teams some more
204 lines
10 KiB
Python
204 lines
10 KiB
Python
"""
|
|
Tests for litellm/proxy/client/cli/commands/auth.py
|
|
|
|
This module tests the auth commands and their associated functionality.
|
|
"""
|
|
|
|
import pytest
|
|
import requests
|
|
from unittest.mock import AsyncMock, patch, Mock, call
|
|
from litellm.proxy.client.cli.commands.auth import _normalize_teams, _poll_for_ready_data, _poll_for_authentication
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_normalize_teams_teams_only():
|
|
"""Test normalize teams helper function"""
|
|
teams = ["1", "2", "3"]
|
|
team_details = []
|
|
result = _normalize_teams(teams, team_details)
|
|
assert result == [{"team_id": "1", "team_alias": None}, {"team_id": "2", "team_alias": None}, {"team_id": "3", "team_alias": None}]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_normalize_teams_with_details_no_aliases():
|
|
"""Test normalize teams helper function"""
|
|
teams = ["4", "5", "6"]
|
|
team_details = [{"team_id": "1"}, {"team_id": "2"}, {"team_id": "3"}]
|
|
result = _normalize_teams(teams, team_details)
|
|
assert result == [{"team_id": "1", "team_alias": None}, {"team_id": "2", "team_alias": None}, {"team_id": "3", "team_alias": None}]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_normalize_teams_with_details_with_aliases():
|
|
"""Test normalize teams helper function"""
|
|
teams = ["4", "5", "6"]
|
|
team_details = [{"team_id": "1", "team_alias": "A"}, {"team_id": "2", "team_alias": "B"}, {"team_id": "3", "team_alias": "C"}]
|
|
result = _normalize_teams(teams, team_details)
|
|
assert result == [{"team_id": "1", "team_alias": "A"}, {"team_id": "2", "team_alias": "B"}, {"team_id": "3", "team_alias": "C"}]
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth.requests.get", side_effect=[Mock(status_code=404)])
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
@patch("litellm.proxy.client.cli.commands.auth.time.sleep")
|
|
async def test_poll_for_ready_404(sleep_mock, click_mock, request_mock):
|
|
"""Test poll_for_ready function"""
|
|
actual = _poll_for_ready_data("https://litellm.com", poll_interval=1, total_timeout=1, request_timeout=42)
|
|
assert actual is None
|
|
click_mock.assert_called_once_with("Polling error: HTTP 404")
|
|
request_mock.assert_called_once_with("https://litellm.com", timeout=42)
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth.requests.get", side_effect=[Mock(status_code=200, json=Mock(return_value={"status": "ready","json": "data"}))])
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
@patch("litellm.proxy.client.cli.commands.auth.time.sleep")
|
|
async def test_poll_for_ready_200_ready(sleep_mock, click_mock, request_mock):
|
|
"""Test poll_for_ready function"""
|
|
actual = _poll_for_ready_data("https://litellm.com", poll_interval=1, total_timeout=1, request_timeout=42)
|
|
assert actual == {"status": "ready", "json": "data"}
|
|
click_mock.assert_not_called()
|
|
request_mock.assert_called_once_with("https://litellm.com", timeout=42)
|
|
sleep_mock.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth.requests.get", side_effect=[Mock(status_code=200, json=Mock(return_value={"status": "pending","json": "data"})), Mock(status_code=200, json=Mock(return_value={"status": "ready","json": "data"}))])
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
@patch("litellm.proxy.client.cli.commands.auth.time.sleep")
|
|
async def test_poll_for_ready_single_pending(sleep_mock, click_mock, request_mock):
|
|
"""Test poll_for_ready function"""
|
|
actual = _poll_for_ready_data("https://litellm.com", poll_interval=1, total_timeout=2, request_timeout=42)
|
|
assert actual == {"status": "ready", "json": "data"}
|
|
click_mock.assert_not_called()
|
|
request_mock.assert_has_calls([
|
|
call("https://litellm.com", timeout=42),
|
|
call("https://litellm.com", timeout=42)
|
|
])
|
|
sleep_mock.assert_called_once_with(1)
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth.requests.get", side_effect=[Mock(status_code=200, json=Mock(return_value={"status": "pending","json": "data"})), Mock(status_code=200, json=Mock(return_value={"status": "pending","json": "data"}))])
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
@patch("litellm.proxy.client.cli.commands.auth.time.sleep")
|
|
async def test_poll_for_ready_pending(sleep_mock, click_mock, request_mock):
|
|
"""Test poll_for_ready function"""
|
|
actual = _poll_for_ready_data("https://litellm.com", poll_interval=1, total_timeout=2, request_timeout=42, pending_message="Pending message", pending_log_every=1)
|
|
assert actual is None
|
|
click_mock.assert_has_calls([
|
|
call("Pending message"),
|
|
call("Pending message")
|
|
])
|
|
request_mock.assert_has_calls([
|
|
call("https://litellm.com", timeout=42),
|
|
call("https://litellm.com", timeout=42)
|
|
])
|
|
sleep_mock.assert_has_calls([
|
|
call(1),
|
|
call(1)
|
|
])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth.requests.get", side_effect=[requests.RequestException("ERROR"),
|
|
requests.RequestException("ERROR")])
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
@patch("litellm.proxy.client.cli.commands.auth.time.sleep")
|
|
async def test_poll_for_ready_connection_failure(sleep_mock, click_mock, request_mock):
|
|
"""Test poll_for_ready function"""
|
|
actual = _poll_for_ready_data("https://litellm.com", poll_interval=1, total_timeout=2, request_timeout=42)
|
|
assert actual is None
|
|
click_mock.assert_called_once_with("Connection error (will retry): ERROR")
|
|
request_mock.assert_has_calls([
|
|
call("https://litellm.com", timeout=42),
|
|
])
|
|
sleep_mock.assert_has_calls([
|
|
call(1),
|
|
call(1)
|
|
])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth._handle_team_selection_during_polling")
|
|
@patch("litellm.proxy.client.cli.commands.auth._poll_for_ready_data", return_value=None)
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
async def test_poll_for_authentication_no_data(click_mock, poll_mock, handle_mock):
|
|
"""Test poll_for_authentication function"""
|
|
actual = _poll_for_authentication("https://litellm.com", "key-123")
|
|
assert actual is None
|
|
poll_mock.assert_called_once_with(
|
|
"https://litellm.com/sso/cli/poll/key-123",
|
|
pending_message="Still waiting for authentication...",
|
|
)
|
|
handle_mock.assert_not_called()
|
|
click_mock.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth._handle_team_selection_during_polling")
|
|
@patch("litellm.proxy.client.cli.commands.auth._poll_for_ready_data", return_value={"requires_team_selection": True, "teams": [], "team_details": []})
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
async def test_poll_for_authentication_no_teams(click_mock, poll_mock, handle_mock):
|
|
"""Test poll_for_authentication function"""
|
|
actual = _poll_for_authentication("https://litellm.com", "key-123")
|
|
assert actual is None
|
|
poll_mock.assert_called_once_with(
|
|
"https://litellm.com/sso/cli/poll/key-123",
|
|
pending_message="Still waiting for authentication...",
|
|
)
|
|
handle_mock.assert_not_called()
|
|
click_mock.assert_called_once()
|
|
assert "No teams available for selection." in click_mock.call_args[0][0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth._handle_team_selection_during_polling", return_value="jwt-123")
|
|
@patch("litellm.proxy.client.cli.commands.auth._poll_for_ready_data", return_value={"requires_team_selection": True, "teams": [1, 2], "user_id": "user-123"})
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
async def test_poll_for_authentication_team_selection_success(click_mock, poll_mock, handle_mock):
|
|
"""Test poll_for_authentication function"""
|
|
actual = _poll_for_authentication("https://litellm.com", "key-123")
|
|
assert actual == {"api_key": "jwt-123", "user_id": "user-123", "teams": [1, 2], "team_id": None}
|
|
poll_mock.assert_called_once_with(
|
|
"https://litellm.com/sso/cli/poll/key-123",
|
|
pending_message="Still waiting for authentication...",
|
|
)
|
|
handle_mock.assert_called_once_with(
|
|
base_url="https://litellm.com",
|
|
key_id="key-123",
|
|
teams=[{"team_id": "1", "team_alias": None}, {"team_id": "2", "team_alias": None}],
|
|
)
|
|
click_mock.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth._handle_team_selection_during_polling", return_value=None)
|
|
@patch("litellm.proxy.client.cli.commands.auth._poll_for_ready_data", return_value={"requires_team_selection": True, "teams": ["team-1"], "user_id": "user-123"})
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
async def test_poll_for_authentication_team_selection_cancelled(click_mock, poll_mock, handle_mock):
|
|
"""Test poll_for_authentication function"""
|
|
actual = _poll_for_authentication("https://litellm.com", "key-123")
|
|
assert actual is None
|
|
poll_mock.assert_called_once_with(
|
|
"https://litellm.com/sso/cli/poll/key-123",
|
|
pending_message="Still waiting for authentication...",
|
|
)
|
|
handle_mock.assert_called_once_with(
|
|
base_url="https://litellm.com",
|
|
key_id="key-123",
|
|
teams=[{"team_id": "team-1", "team_alias": None}],
|
|
)
|
|
click_mock.assert_called_once()
|
|
assert "Team selection cancelled" in click_mock.call_args[0][0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("litellm.proxy.client.cli.commands.auth._handle_team_selection_during_polling")
|
|
@patch("litellm.proxy.client.cli.commands.auth._poll_for_ready_data", return_value={"key": "jwt-456", "user_id": "user-456", "teams": ["team-1"], "team_id": "team-1"})
|
|
@patch("litellm.proxy.client.cli.commands.auth.click.echo")
|
|
async def test_poll_for_authentication_auto_assigned_team(click_mock, poll_mock, handle_mock):
|
|
"""Test poll_for_authentication function"""
|
|
actual = _poll_for_authentication("https://litellm.com", "key-123")
|
|
assert actual == {"api_key": "jwt-456", "user_id": "user-456", "teams": ["team-1"], "team_id": "team-1"}
|
|
poll_mock.assert_called_once_with(
|
|
"https://litellm.com/sso/cli/poll/key-123",
|
|
pending_message="Still waiting for authentication...",
|
|
)
|
|
handle_mock.assert_not_called()
|
|
click_mock.assert_called_once()
|
|
assert "Automatically assigned to team: team-1" in click_mock.call_args[0][0]
|