vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

server: fix disconnection logic in test (before post response headers)

Browse source
This commit is contained in:
ochafik 2024-10-05 00:44:13 +01:00
parent 6f693f14b0
commit 52c5a6244f
1 changed files with 6 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

14
examples/server/tests/features/steps/steps.py
View file

@ -1012,6 +1012,9 @@ async def request_completion(prompt,
headers['Authorization'] = f'Bearer {user_api_key}' headers['Authorization'] = f'Bearer {user_api_key}'
async with aiohttp.ClientSession(timeout=DEFAULT_TIMEOUT_SECONDS) as session: async with aiohttp.ClientSession(timeout=DEFAULT_TIMEOUT_SECONDS) as session:
if disconnect_after_millis is not None:
await asyncio.sleep(disconnect_after_millis / 1000.0)
return 0
async with session.post(f'{base_url}/completion', async with session.post(f'{base_url}/completion',
json={ json={
"input_prefix": prompt_prefix, "input_prefix": prompt_prefix,
@ -1025,10 +1028,6 @@ async def request_completion(prompt,
"n_probs": 2, "n_probs": 2,
}, },
headers=headers) as response: headers=headers) as response:
if disconnect_after_millis is not None:
await asyncio.sleep(disconnect_after_millis / 1000.0)
return 0
if expect_api_error is None or not expect_api_error: if expect_api_error is None or not expect_api_error:
assert response.status == 200 assert response.status == 200
assert response.headers['Access-Control-Allow-Origin'] == origin assert response.headers['Access-Control-Allow-Origin'] == origin
@ -1088,13 +1087,12 @@ async def oai_chat_completions(user_prompt,
origin = 'llama.cpp' origin = 'llama.cpp'
headers = {'Authorization': f'Bearer {user_api_key}', 'Origin': origin} headers = {'Authorization': f'Bearer {user_api_key}', 'Origin': origin}
async with aiohttp.ClientSession(timeout=DEFAULT_TIMEOUT_SECONDS) as session: async with aiohttp.ClientSession(timeout=DEFAULT_TIMEOUT_SECONDS) as session:
if disconnect_after_millis is not None:
await asyncio.sleep(disconnect_after_millis / 1000.0)
return 0
async with session.post(f'{base_url}{base_path}', async with session.post(f'{base_url}{base_path}',
json=payload, json=payload,
headers=headers) as response: headers=headers) as response:
if disconnect_after_millis is not None:
await asyncio.sleep(disconnect_after_millis / 1000.0)
return 0
if enable_streaming: if enable_streaming:
assert response.status == 200 assert response.status == 200
assert response.headers['Access-Control-Allow-Origin'] == origin assert response.headers['Access-Control-Allow-Origin'] == origin