Skip to content

Commit aaa8f04

Browse files
authored
Reduce API cross-section for huggingface in test (#3042)
1 parent 41aa99b commit aaa8f04

File tree

1 file changed

+16
-40
lines changed

1 file changed

+16
-40
lines changed

tests/integrations/huggingface_hub/test_huggingface_hub.py

Lines changed: 16 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,8 @@
11
import itertools
2-
import json
32

43
import pytest
54
from huggingface_hub import (
65
InferenceClient,
7-
TextGenerationOutput,
8-
TextGenerationOutputDetails,
9-
TextGenerationStreamOutput,
10-
TextGenerationOutputToken,
11-
TextGenerationStreamDetails,
126
)
137
from huggingface_hub.errors import OverloadedError
148

@@ -35,19 +29,15 @@ def test_nonstreaming_chat_completion(
3529
client = InferenceClient("some-model")
3630
if details_arg:
3731
client.post = mock.Mock(
38-
return_value=json.dumps(
39-
[
40-
TextGenerationOutput(
41-
generated_text="the model response",
42-
details=TextGenerationOutputDetails(
43-
finish_reason="TextGenerationFinishReason",
44-
generated_tokens=10,
45-
prefill=[],
46-
tokens=[], # not needed for integration
47-
),
48-
)
49-
]
50-
).encode("utf-8")
32+
return_value=b"""[{
33+
"generated_text": "the model response",
34+
"details": {
35+
"finish_reason": "length",
36+
"generated_tokens": 10,
37+
"prefill": [],
38+
"tokens": []
39+
}
40+
}]"""
5141
)
5242
else:
5343
client.post = mock.Mock(
@@ -96,27 +86,13 @@ def test_streaming_chat_completion(
9686
client = InferenceClient("some-model")
9787
client.post = mock.Mock(
9888
return_value=[
99-
b"data:"
100-
+ json.dumps(
101-
TextGenerationStreamOutput(
102-
token=TextGenerationOutputToken(
103-
id=1, special=False, text="the model "
104-
),
105-
),
106-
).encode("utf-8"),
107-
b"data:"
108-
+ json.dumps(
109-
TextGenerationStreamOutput(
110-
token=TextGenerationOutputToken(
111-
id=2, special=False, text="response"
112-
),
113-
details=TextGenerationStreamDetails(
114-
finish_reason="length",
115-
generated_tokens=10,
116-
seed=0,
117-
),
118-
)
119-
).encode("utf-8"),
89+
b"""data:{
90+
"token":{"id":1, "special": false, "text": "the model "}
91+
}""",
92+
b"""data:{
93+
"token":{"id":2, "special": false, "text": "response"},
94+
"details":{"finish_reason": "length", "generated_tokens": 10, "seed": 0}
95+
}""",
12096
]
12197
)
12298
with start_transaction(name="huggingface_hub tx"):

0 commit comments

Comments
 (0)