litellm

1
# What this tests?
2
## Tests /health + /routes endpoints.
3

4
import pytest
5
import asyncio
6
import aiohttp
7

8

9
async def health(session, call_key):
10
    url = "http://0.0.0.0:4000/health"
11
    headers = {
12
        "Authorization": f"Bearer {call_key}",
13
        "Content-Type": "application/json",
14
    }
15

16
    async with session.get(url, headers=headers) as response:
17
        status = response.status
18
        response_text = await response.text()
19

20
        print(f"Response (Status code: {status}):")
21
        print(response_text)
22
        print()
23

24
        if status != 200:
25
            raise Exception(f"Request did not return a 200 status code: {status}")
26

27
        return await response.json()
28

29

30
async def generate_key(session):
31
    url = "http://0.0.0.0:4000/key/generate"
32
    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
33
    data = {
34
        "models": ["gpt-4", "text-embedding-ada-002", "dall-e-2"],
35
        "duration": None,
36
    }
37

38
    async with session.post(url, headers=headers, json=data) as response:
39
        status = response.status
40
        response_text = await response.text()
41

42
        print(response_text)
43
        print()
44

45
        if status != 200:
46
            raise Exception(f"Request did not return a 200 status code: {status}")
47
        return await response.json()
48

49

50
@pytest.mark.asyncio
51
async def test_health():
52
    """
53
    - Call /health
54
    """
55
    async with aiohttp.ClientSession() as session:
56
        # as admin #
57
        all_healthy_models = await health(session=session, call_key="sk-1234")
58
        total_model_count = (
59
            all_healthy_models["healthy_count"] + all_healthy_models["unhealthy_count"]
60
        )
61
        assert total_model_count > 0
62

63

64
@pytest.mark.asyncio
65
async def test_health_readiness():
66
    """
67
    Check if 200
68
    """
69
    async with aiohttp.ClientSession() as session:
70
        url = "http://0.0.0.0:4000/health/readiness"
71
        async with session.get(url) as response:
72
            status = response.status
73
            response_json = await response.json()
74

75
            print(response_json)
76
            assert "litellm_version" in response_json
77
            assert "status" in response_json
78

79
            if status != 200:
80
                raise Exception(f"Request did not return a 200 status code: {status}")
81

82

83
@pytest.mark.asyncio
84
async def test_health_liveliness():
85
    """
86
    Check if 200
87
    """
88
    async with aiohttp.ClientSession() as session:
89
        url = "http://0.0.0.0:4000/health/liveliness"
90
        async with session.get(url) as response:
91
            status = response.status
92
            response_text = await response.text()
93

94
            print(response_text)
95
            print()
96

97
            if status != 200:
98
                raise Exception(f"Request did not return a 200 status code: {status}")
99

100

101
@pytest.mark.asyncio
102
async def test_routes():
103
    """
104
    Check if 200
105
    """
106
    async with aiohttp.ClientSession() as session:
107
        url = "http://0.0.0.0:4000/routes"
108
        async with session.get(url) as response:
109
            status = response.status
110
            response_text = await response.text()
111

112
            print(response_text)
113
            print()
114

115
            if status != 200:
116
                raise Exception(f"Request did not return a 200 status code: {status}")
117