test: add test for record and chunk

vfa-khuongdv · May 28, 2024 · 1fa1abd · 1fa1abd
1 parent 1b68146
commit 1fa1abd
Show file tree

Hide file tree

Showing 5 changed files with 112 additions and 107 deletions.
diff --git a/backend/app/routes/inference/chat_completion.py b/backend/app/routes/inference/chat_completion.py
@@ -57,19 +57,21 @@ async def api_chat_completion(
             if not model.allow_streaming():
                 raise_request_validation_error(f"Model {model.model_id} does not support streaming.")
 
-            async def generator():
-                async for chunk_dict in await stream_chat_completion(
-                    model=model,
-                    messages=messages,
-                    configs=data.configs,
-                    function_call=data.function_call,
-                    functions=functions,
-                ):
+            async def generator(sse_chunk_dicts):
+                async for chunk_dict in sse_chunk_dicts:
                     yield f"data: {json.dumps(chunk_dict)}\n\n"
                 yield SSE_DONE_MSG
 
+            sse_chunk_dicts = await stream_chat_completion(
+                model=model,
+                messages=messages,
+                configs=data.configs,
+                function_call=data.function_call,
+                functions=functions,
+            )
+
             return StreamingResponse(
-                generator(),
+                generator(sse_chunk_dicts),
                 media_type="text/event-stream",
             )
 

diff --git a/backend/requirements.txt b/backend/requirements.txt
@@ -31,4 +31,5 @@ markdown==3.6
 pytest==7.4.3
 pytest-asyncio==0.21.1
 pytest-ordering==0.6
-pytest-assume==2.4.3
+pytest-assume==2.4.3
+pytest-rerunfailures==13.0
diff --git a/backend/run_api_test.sh b/backend/run_api_test.sh
@@ -5,5 +5,5 @@ parent_dir="$(dirname "$(pwd)")"
 export PYTHONPATH="${PYTHONPATH}:${parent_dir}"
 
 echo "Starting tests..."
-pytest   ./tests/services_tests  -m "api_test"  -q  --tb=no
+pytest   ./tests/services_tests  -m "api_test"  -q  --tb=no  --reruns 3 --reruns-delay 2
 echo "Tests completed."
diff --git a/backend/run_web_test.sh b/backend/run_web_test.sh
@@ -4,5 +4,5 @@ set -e
 parent_dir="$(dirname "$(pwd)")"
 export PYTHONPATH="${PYTHONPATH}:${parent_dir}"
 echo "Starting tests..."
-pytest   ./tests/services_tests  -m "web_test or api_test"
+pytest   ./tests/services_tests  -m "web_test or api_test"  -q  --tb=no  --reruns 3 --reruns-delay 2
 echo "Tests completed."