From 591bd9fcdd3016145191ecfb52b271da69e3a85b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 5 Apr 2024 12:49:29 -0700 Subject: [PATCH] test(test_openai_endpoints.py): add response header test - make sure it's always <4kb --- tests/test_openai_endpoints.py | 167 ++++++++++++++++++++------------- 1 file changed, 100 insertions(+), 67 deletions(-) diff --git a/tests/test_openai_endpoints.py b/tests/test_openai_endpoints.py index 9535b4842..2137b2665 100644 --- a/tests/test_openai_endpoints.py +++ b/tests/test_openai_endpoints.py @@ -6,6 +6,14 @@ import aiohttp, openai from openai import OpenAI +def response_header_check(response): + """ + - assert if response headers < 4kb (nginx limit). + """ + headers_size = sum(len(k) + len(v) for k, v in response.raw_headers) + assert headers_size < 4096, "Response headers exceed the 4kb limit" + + async def generate_key(session): url = "http://0.0.0.0:4000/key/generate" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} @@ -23,6 +31,11 @@ async def generate_key(session): if status != 200: raise Exception(f"Request did not return a 200 status code: {status}") + + response_header_check( + response + ) # calling the function to check response headers + return await response.json() @@ -43,6 +56,10 @@ async def new_user(session): if status != 200: raise Exception(f"Request did not return a 200 status code: {status}") + + response_header_check( + response + ) # calling the function to check response headers return await response.json() @@ -69,9 +86,92 @@ async def chat_completion(session, key): if status != 200: raise Exception(f"Request did not return a 200 status code: {status}") + + response_header_check( + response + ) # calling the function to check response headers + return await response.json() +async def completion(session, key): + url = "http://0.0.0.0:4000/completions" + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + } + data = {"model": "gpt-4", "prompt": "Hello!"} + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + + if status != 200: + raise Exception(f"Request did not return a 200 status code: {status}") + + response_header_check( + response + ) # calling the function to check response headers + + response = await response.json() + + return response + + +async def embeddings(session, key): + url = "http://0.0.0.0:4000/embeddings" + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + } + data = { + "model": "text-embedding-ada-002", + "input": ["hello world"], + } + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + response_text = await response.text() + + print(response_text) + + if status != 200: + raise Exception(f"Request did not return a 200 status code: {status}") + + response_header_check( + response + ) # calling the function to check response headers + + +async def image_generation(session, key): + url = "http://0.0.0.0:4000/images/generations" + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + } + data = { + "model": "dall-e-2", + "prompt": "A cute baby sea otter", + } + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + response_text = await response.text() + + print(response_text) + print() + + if status != 200: + if ( + "Connection error" in response_text + ): # OpenAI endpoint returns a connection error + return + raise Exception(f"Request did not return a 200 status code: {status}") + + response_header_check( + response + ) # calling the function to check response headers + + @pytest.mark.asyncio async def test_chat_completion(): """ @@ -105,25 +205,6 @@ async def test_chat_completion_old_key(): await chat_completion(session=session, key=key) -async def completion(session, key): - url = "http://0.0.0.0:4000/completions" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - data = {"model": "gpt-4", "prompt": "Hello!"} - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - - if status != 200: - raise Exception(f"Request did not return a 200 status code: {status}") - - response = await response.json() - - return response - - @pytest.mark.asyncio async def test_completion(): """ @@ -151,28 +232,6 @@ async def test_completion(): ) -async def embeddings(session, key): - url = "http://0.0.0.0:4000/embeddings" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - data = { - "model": "text-embedding-ada-002", - "input": ["hello world"], - } - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - print(response_text) - print() - - if status != 200: - raise Exception(f"Request did not return a 200 status code: {status}") - - @pytest.mark.asyncio async def test_embeddings(): """ @@ -190,32 +249,6 @@ async def test_embeddings(): await embeddings(session=session, key=key_2) -async def image_generation(session, key): - url = "http://0.0.0.0:4000/images/generations" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - data = { - "model": "dall-e-2", - "prompt": "A cute baby sea otter", - } - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - print(response_text) - print() - - if status != 200: - if ( - "Connection error" in response_text - ): # OpenAI endpoint returns a connection error - return - raise Exception(f"Request did not return a 200 status code: {status}") - - @pytest.mark.asyncio async def test_image_generation(): """