FilesExpand file tree

 main

/

test_client.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

116 lines (92 loc) · 3.24 KB

 main

/

test_client.py

Copy path

Top

File metadata and controls

• Code
• Blame

116 lines (92 loc) · 3.24 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

import os

from unittest import mock

import httpx

import pytest

import respx

@pytest.mark.asyncio

async def test_authorization_when_setting_environ_after_import():

import replicate

router = respx.Router()

router.route(

method="GET",

url="https://api.replicate.com/",

headers={"Authorization": "Bearer test-set-after-import"},

).mock(

return_value=httpx.Response(

200,

json={},

)

)

token = "test-set-after-import" # noqa: S105

with mock.patch.dict(

os.environ,

{"REPLICATE_API_TOKEN": token},

):

client = replicate.Client(transport=httpx.MockTransport(router.handler))

resp = client._request("GET", "/")

assert resp.status_code == 200

@pytest.mark.asyncio

async def test_client_error_handling():

import replicate

from replicate.exceptions import ReplicateError

router = respx.Router()

router.route(

method="GET",

url="https://api.replicate.com/",

headers={"Authorization": "Bearer test-client-error"},

).mock(

return_value=httpx.Response(

400,

json={"detail": "Client error occurred"},

)

)

token = "test-client-error" # noqa: S105

with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": token}):

client = replicate.Client(transport=httpx.MockTransport(router.handler))

with pytest.raises(ReplicateError) as exc_info:

client._request("GET", "/")

assert "status: 400" in str(exc_info.value)

assert "detail: Client error occurred" in str(exc_info.value)

@pytest.mark.asyncio

async def test_server_error_handling():

import replicate

from replicate.exceptions import ReplicateError

router = respx.Router()

router.route(

method="GET",

url="https://api.replicate.com/",

headers={"Authorization": "Bearer test-server-error"},

).mock(

return_value=httpx.Response(

500,

json={"detail": "Server error occurred"},

)

)

token = "test-server-error" # noqa: S105

with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": token}):

client = replicate.Client(transport=httpx.MockTransport(router.handler))

with pytest.raises(ReplicateError) as exc_info:

client._request("GET", "/")

assert "status: 500" in str(exc_info.value)

assert "detail: Server error occurred" in str(exc_info.value)

def test_custom_headers_are_applied():

import replicate

from replicate.exceptions import ReplicateError

custom_headers = {"User-Agent": "my-custom-user-agent/1.0"}

def mock_send(request):

assert "User-Agent" in request.headers, "Custom header not found in request"

assert (

request.headers["User-Agent"] == "my-custom-user-agent/1.0"

), "Custom header value is incorrect"

return httpx.Response(401, json={})

mock_send_wrapper = mock.Mock(side_effect=mock_send)

client = replicate.Client(

api_token="dummy_token",

headers=custom_headers,

transport=httpx.MockTransport(mock_send_wrapper),

)

try:

client.accounts.current()

except ReplicateError:

pass

mock_send_wrapper.assert_called_once()