Skip to content

Commit a8e5295

Browse files
authored
[NeuralChat] Fix tgi endpoint in test (intel#1388)
1 parent 8797cfe commit a8e5295

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

intel_extension_for_transformers/neural_chat/tests/nightly/serving/test_tgi_serving.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ def test_tgi_root(self, mock_text_generation):
4444
"parameters": {"max_new_tokens":10},
4545
"stream": False
4646
}
47+
router.set_tgi_endpoint("/v1/tgi")
4748
response = client.post("/v1/tgi", json=request_data)
4849

4950
mock_text_generation.assert_called_once_with(
@@ -67,6 +68,7 @@ def test_tgi_generate(self, mock_text_generation):
6768
"inputs": "Test text generation inputs.",
6869
"parameters": {"max_new_tokens":10}
6970
}
71+
router.set_tgi_endpoint("/v1/tgi/generate")
7072
response = client.post("/v1/tgi/generate", json=request_data)
7173

7274
mock_text_generation.assert_called_once_with(
@@ -93,6 +95,7 @@ def mock_generator():
9395
"inputs": "Test text generation inputs.",
9496
"parameters": {"max_new_tokens":10}
9597
}
98+
router.set_tgi_endpoint("/v1/tgi/generate_stream")
9699
response = client.post("/v1/tgi/generate_stream", json=request_data)
97100

98101
self.assertEqual(response.status_code, 200)

0 commit comments

Comments
 (0)