text-generation-inference

1
from text_generation.errors import (
2
    parse_error,
3
    GenerationError,
4
    IncompleteGenerationError,
5
    OverloadedError,
6
    ValidationError,
7
    BadRequestError,
8
    ShardNotReadyError,
9
    ShardTimeoutError,
10
    NotFoundError,
11
    RateLimitExceededError,
12
    UnknownError,
13
)
14

15

16
def test_generation_error():
17
    payload = {"error_type": "generation", "error": "test"}
18
    assert isinstance(parse_error(400, payload), GenerationError)
19

20

21
def test_incomplete_generation_error():
22
    payload = {"error_type": "incomplete_generation", "error": "test"}
23
    assert isinstance(parse_error(400, payload), IncompleteGenerationError)
24

25

26
def test_overloaded_error():
27
    payload = {"error_type": "overloaded", "error": "test"}
28
    assert isinstance(parse_error(400, payload), OverloadedError)
29

30

31
def test_validation_error():
32
    payload = {"error_type": "validation", "error": "test"}
33
    assert isinstance(parse_error(400, payload), ValidationError)
34

35

36
def test_bad_request_error():
37
    payload = {"error": "test"}
38
    assert isinstance(parse_error(400, payload), BadRequestError)
39

40

41
def test_shard_not_ready_error():
42
    payload = {"error": "test"}
43
    assert isinstance(parse_error(403, payload), ShardNotReadyError)
44
    assert isinstance(parse_error(424, payload), ShardNotReadyError)
45

46

47
def test_shard_timeout_error():
48
    payload = {"error": "test"}
49
    assert isinstance(parse_error(504, payload), ShardTimeoutError)
50

51

52
def test_not_found_error():
53
    payload = {"error": "test"}
54
    assert isinstance(parse_error(404, payload), NotFoundError)
55

56

57
def test_rate_limit_exceeded_error():
58
    payload = {"error": "test"}
59
    assert isinstance(parse_error(429, payload), RateLimitExceededError)
60

61

62
def test_unknown_error():
63
    payload = {"error": "test"}
64
    assert isinstance(parse_error(500, payload), UnknownError)
65