File size: 2,673 Bytes
447ebeb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
from unittest.mock import patch

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_users_client():
    with patch(
        "litellm.proxy.client.cli.commands.users.UsersManagementClient"
    ) as MockClient:
        yield MockClient


def test_users_list(cli_runner, mock_users_client):
    mock_users_client.return_value.list_users.return_value = [
        {
            "user_id": "u1",
            "user_email": "[email protected]",
            "user_role": "internal_user",
            "teams": ["t1"],
        },
        {
            "user_id": "u2",
            "user_email": "[email protected]",
            "user_role": "proxy_admin",
            "teams": ["t2", "t3"],
        },
    ]
    result = cli_runner.invoke(cli, ["users", "list"])
    assert result.exit_code == 0
    assert "u1" in result.output
    assert "[email protected]" in result.output
    assert "proxy_admin" in result.output
    assert "t3" in result.output
    mock_users_client.return_value.list_users.assert_called_once()


def test_users_get(cli_runner, mock_users_client):
    mock_users_client.return_value.get_user.return_value = {
        "user_id": "u1",
        "user_email": "[email protected]",
    }
    result = cli_runner.invoke(cli, ["users", "get", "--id", "u1"])
    assert result.exit_code == 0
    assert '"user_id": "u1"' in result.output
    assert '"user_email": "[email protected]"' in result.output
    mock_users_client.return_value.get_user.assert_called_once_with(user_id="u1")


def test_users_create(cli_runner, mock_users_client):
    mock_users_client.return_value.create_user.return_value = {
        "user_id": "u1",
        "user_email": "[email protected]",
    }
    result = cli_runner.invoke(
        cli, ["users", "create", "--email", "[email protected]", "--role", "internal_user"]
    )
    assert result.exit_code == 0
    assert '"user_id": "u1"' in result.output
    assert '"user_email": "[email protected]"' in result.output
    mock_users_client.return_value.create_user.assert_called_once()


def test_users_delete(cli_runner, mock_users_client):
    mock_users_client.return_value.delete_user.return_value = {"deleted": 1}
    result = cli_runner.invoke(cli, ["users", "delete", "u1", "u2"])
    assert result.exit_code == 0
    assert '"deleted": 1' in result.output
    mock_users_client.return_value.delete_user.assert_called_once_with(["u1", "u2"])