server: tests: move all requests call to asyncio

This commit is contained in:
Pierrick HYMBERT 2024-02-23 14:44:12 +01:00
parent e10b83a217
commit 4d27466ca5

View file

@ -10,7 +10,6 @@ from re import RegexFlag
import aiohttp import aiohttp
import openai import openai
import requests
from behave import step from behave import step
from behave.api.async_step import async_run_until_complete from behave.api.async_step import async_run_until_complete
@ -75,6 +74,7 @@ def step_server_n_predict(context, n_predict):
def step_server_continuous_batching(context): def step_server_continuous_batching(context):
context.server_continuous_batching = True context.server_continuous_batching = True
@step(u'embeddings extraction') @step(u'embeddings extraction')
def step_server_embeddings(context): def step_server_embeddings(context):
context.server_embeddings = True context.server_embeddings = True
@ -306,12 +306,16 @@ async def all_prompts_are_predicted(context, expected_predicted_n=None):
@step(u'embeddings are computed for') @step(u'embeddings are computed for')
def step_compute_embedding(context): @async_run_until_complete
response = requests.post(f'{context.base_url}/embedding', json={ async def step_compute_embedding(context):
"content": context.text, async with aiohttp.ClientSession() as session:
}) async with session.post(f'{context.base_url}/embedding',
assert response.status_code == 200 json={
context.embeddings = response.json()['embedding'] "content": context.text,
}) as response:
assert response.status == 200
response_json = await response.json()
context.embeddings = response_json['embedding']
@step(u'embeddings are generated') @step(u'embeddings are generated')
@ -338,32 +342,42 @@ def step_oai_compute_embedding(context):
@step(u'tokenizing') @step(u'tokenizing')
def step_tokenize(context): @async_run_until_complete
async def step_tokenize(context):
context.tokenized_text = context.text context.tokenized_text = context.text
response = requests.post(f'{context.base_url}/tokenize', json={ async with aiohttp.ClientSession() as session:
"content": context.tokenized_text, async with session.post(f'{context.base_url}/tokenize',
}) json={
assert response.status_code == 200 "content": context.tokenized_text,
context.tokens = response.json()['tokens'] }) as response:
assert response.status == 200
tokenize_json = await response.json()
context.tokens = tokenize_json['tokens']
@step(u'tokens can be detokenize') @step(u'tokens can be detokenize')
def step_detokenize(context): @async_run_until_complete
async def step_detokenize(context):
assert len(context.tokens) > 0 assert len(context.tokens) > 0
response = requests.post(f'{context.base_url}/detokenize', json={ async with aiohttp.ClientSession() as session:
"tokens": context.tokens, async with session.post(f'{context.base_url}/detokenize',
}) json={
assert response.status_code == 200 "tokens": context.tokens,
# SPM tokenizer adds a whitespace prefix: https://github.com/google/sentencepiece/issues/15 }) as response:
assert context.tokenized_text == response.json()['content'].strip() assert response.status == 200
detokenize_json = await response.json()
# SPM tokenizer adds a whitespace prefix: https://github.com/google/sentencepiece/issues/15
assert context.tokenized_text == detokenize_json['content'].strip()
@step(u'an OPTIONS request is sent from {origin}') @step(u'an OPTIONS request is sent from {origin}')
def step_options_request(context, origin): @async_run_until_complete
options_response = requests.options(f'{context.base_url}/v1/chat/completions', async def step_options_request(context, origin):
headers={"Origin": origin}) async with aiohttp.ClientSession() as session:
assert options_response.status_code == 200 async with session.options(f'{context.base_url}/v1/chat/completions',
context.options_response = options_response headers={"Origin": origin}) as response:
assert response.status == 200
context.options_response = response
@step(u'CORS header {cors_header} is set to {cors_header_value}') @step(u'CORS header {cors_header} is set to {cors_header_value}')