chore: add download success e2e test (#1556)

* chore: add download success e2e test

* update

Author: namchuai

.github/workflows/cortex-cpp-quality-gate.yml

@@ -3,11 +3,7 @@ name: CI Quality Gate Cortex CPP
 on:
   pull_request:
     types: [opened, synchronize, reopened, ready_for_review]
-    paths:
-      [
-        "engine/**",
-        ".github/workflows/cortex-cpp-quality-gate.yml"
-      ]
+    paths: ["engine/**", ".github/workflows/cortex-cpp-quality-gate.yml"]
   workflow_dispatch:
 
 env:
@@ -27,25 +23,25 @@ jobs:
             runs-on: "ubuntu-20-04-cuda-12-0"
             cmake-flags: "-DCORTEX_CPP_VERSION=${{github.event.pull_request.head.sha}} -DCMAKE_BUILD_TEST=ON -DCMAKE_TOOLCHAIN_FILE=vcpkg/scripts/buildsystems/vcpkg.cmake"
             build-deps-cmake-flags: ""
-            ccache-dir: ''
+            ccache-dir: ""
           - os: "mac"
             name: "amd64"
             runs-on: "macos-selfhosted-12"
             cmake-flags: "-DCORTEX_CPP_VERSION=${{github.event.pull_request.head.sha}} -DCMAKE_BUILD_TEST=ON -DCMAKE_TOOLCHAIN_FILE=vcpkg/scripts/buildsystems/vcpkg.cmake"
             build-deps-cmake-flags: ""
-            ccache-dir: ''
+            ccache-dir: ""
           - os: "mac"
             name: "arm64"
             runs-on: "macos-silicon"
             cmake-flags: "-DCORTEX_CPP_VERSION=${{github.event.pull_request.head.sha}} -DCMAKE_BUILD_TEST=ON -DMAC_ARM64=ON -DCMAKE_TOOLCHAIN_FILE=vcpkg/scripts/buildsystems/vcpkg.cmake"
             build-deps-cmake-flags: ""
-            ccache-dir: ''
+            ccache-dir: ""
           - os: "windows"
             name: "amd64"
             runs-on: "windows-cuda-12-0"
             cmake-flags: "-DCORTEX_CPP_VERSION=${{github.event.pull_request.head.sha}} -DCMAKE_BUILD_TEST=ON -DBUILD_SHARED_LIBS=OFF -DCMAKE_TOOLCHAIN_FILE=C:/w/cortex.cpp/cortex.cpp/engine/vcpkg/scripts/buildsystems/vcpkg.cmake -DVCPKG_TARGET_TRIPLET=x64-windows-static -DCMAKE_BUILD_TYPE=RELEASE -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache -GNinja"
             build-deps-cmake-flags: "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache -GNinja"
-            ccache-dir: 'C:\Users\ContainerAdministrator\AppData\Local\ccache'          
+            ccache-dir: 'C:\Users\ContainerAdministrator\AppData\Local\ccache'
     steps:
       - name: Clone
         id: checkout
@@ -56,7 +52,7 @@ jobs:
       - name: use python
         uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
+          python-version: "3.10"
 
       - name: Install tools on Linux
         if: runner.os == 'Linux'
@@ -128,8 +124,7 @@ jobs:
           cp build/cortex build/cortex-nightly
           cp build/cortex build/cortex-beta
           python -m pip install --upgrade pip
-          python -m pip install pytest
-          python -m pip install requests
+          python -m pip install -r e2e-test/requirements.txt
           python e2e-test/main.py
           rm build/cortex-nightly
           rm build/cortex-beta
@@ -143,8 +138,7 @@ jobs:
           cp build/cortex.exe build/cortex-nightly.exe
           cp build/cortex.exe build/cortex-beta.exe
           python -m pip install --upgrade pip
-          python -m pip install pytest
-          python -m pip install requests
+          python -m pip install -r e2e-test/requirements.txt
           python e2e-test/main.py
           rm build/cortex-nightly.exe
           rm build/cortex-beta.exe
@@ -155,7 +149,7 @@ jobs:
         run: |
           cd engine
           make pre-package DESTINATION_BINARY_NAME="cortex"
-          
+
       - name: Package
         run: |
           cd engine
@@ -188,4 +182,4 @@ jobs:
         env:
           AWS_ACCESS_KEY_ID: "${{ secrets.MINIO_ACCESS_KEY_ID }}"
           AWS_SECRET_ACCESS_KEY: "${{ secrets.MINIO_SECRET_ACCESS_KEY }}"
-          AWS_DEFAULT_REGION: "${{ secrets.MINIO_REGION }}"
+          AWS_DEFAULT_REGION: "${{ secrets.MINIO_REGION }}"

engine/e2e-test/requirements.txt

@@ -0,0 +1,4 @@
+websockets
+pytest
+pytest-asyncio
+requests

engine/e2e-test/test_api_model_pull_direct_url.py

@@ -1,6 +1,11 @@
 import pytest
 import requests
-from test_runner import start_server, stop_server, run
+from test_runner import (
+    run,
+    start_server,
+    stop_server,
+    wait_for_websocket_download_success_event,
+)
 
 
 class TestApiModelPullDirectUrl:
@@ -30,11 +35,21 @@ def setup_and_teardown(self):
                 "TheBloke:TinyLlama-1.1B-Chat-v0.3-GGUF:tinyllama-1.1b-chat-v0.3.Q2_K.gguf",
             ],
         )
-        stop_server()        
+        stop_server()
 
-    def test_model_pull_with_direct_url_should_be_success(self):
+    @pytest.mark.asyncio
+    async def test_model_pull_with_direct_url_should_be_success(self):
         myobj = {
             "model": "https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/blob/main/tinyllama-1.1b-chat-v0.3.Q2_K.gguf"
         }
         response = requests.post("http://localhost:3928/models/pull", json=myobj)
         assert response.status_code == 200
+        await wait_for_websocket_download_success_event(timeout=None)
+        get_model_response = requests.get(
+            "http://127.0.0.1:3928/models/TheBloke:TinyLlama-1.1B-Chat-v0.3-GGUF:tinyllama-1.1b-chat-v0.3.Q2_K.gguf"
+        )
+        assert get_model_response.status_code == 200
+        assert (
+            get_model_response.json()["model"]
+            == "TheBloke:TinyLlama-1.1B-Chat-v0.3-GGUF:tinyllama-1.1b-chat-v0.3.Q2_K.gguf"
+        )

engine/e2e-test/test_runner.py

@@ -1,3 +1,5 @@
+import asyncio
+import json
 import platform
 import queue
 import select
@@ -6,6 +8,8 @@
 import time
 from typing import List
 
+import websockets
+
 # You might want to change the path of the executable based on your build directory
 executable_windows_path = "build\\cortex.exe"
 executable_unix_path = "build/cortex"
@@ -161,3 +165,25 @@ def enqueue_output(out, queue):
 # Stop the API server
 def stop_server():
     run("Stop server", ["stop"])
+
+
+async def wait_for_websocket_download_success_event(timeout: float = 30):
+    if platform.system() == 'Windows':
+        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+    async with websockets.connect("ws://127.0.0.1:3928/events") as websocket:
+        try:
+            # Using wait_for instead of timeout context manager
+            async def receive_until_success():
+                while True:
+                    message = await websocket.recv()
+                    try:
+                        event = json.loads(message)
+                        if event.get("type") == "DownloadSuccess":
+                            return event
+                    except json.JSONDecodeError:
+                        continue
+
+            return await asyncio.wait_for(receive_until_success(), timeout)
+
+        except asyncio.TimeoutError:
+            raise TimeoutError("Timeout waiting for DownloadSuccess event")