Merge pull request #28 from hv0905/upload-test

Add max queue length config, add upload tests.

app/Controllers/admin.py

@@ -108,7 +108,7 @@ async def upload_image(image_file: Annotated[UploadFile, File(description="The i
     try:
         image = Image.open(BytesIO(img_bytes))
     except UnidentifiedImageError as ex:
-        raise HTTPException(400, "Cannot open the image file.") from ex
+        raise HTTPException(422, "Cannot open the image file.") from ex
 
     image_data = ImageData(id=img_id,
                            url=model.url,

app/Services/upload_service.py

@@ -9,6 +9,7 @@
 from app.Services.index_service import IndexService
 from app.Services.storage import StorageService
 from app.Services.vector_db_context import VectorDbContext
+from app.config import config
 
 
 class UploadService:
@@ -17,7 +18,7 @@ def __init__(self, storage_service: StorageService, db_context: VectorDbContext,
         self._db_context = db_context
         self._index_service = index_service
 
-        self._queue = asyncio.Queue(200)
+        self._queue = asyncio.Queue(config.admin_index_queue_max_length)
         self._upload_worker_task = asyncio.create_task(self._upload_worker())
 
         self._processed_count = 0

app/config.py

@@ -87,6 +87,7 @@ class Config(BaseSettings):
     cors_origins: set[str] = {'*'}
     admin_api_enable: bool = False
     admin_token: str = ''
+    admin_index_queue_max_length: int = 200
 
     access_protected: bool = False
     access_token: str = ''

config/default.env

@@ -71,6 +71,8 @@
 # APP_ADMIN_API_ENABLE=False
 # Uncomment the line below if you enabled admin API. Use this token to access admin API. For security reasons, the admin token is always required if you want to use admin API.
 # APP_ADMIN_TOKEN="your-super-secret-admin-token"
+# Max length of the upload queue for admin API, higher value means more indexing requests can be queued but also means more memory usage. Upload requests will be blocked when the queue is full.
+# APP_ADMIN_INDEX_QUEUE_MAX_LENGTH=200
 
 
 # ------

readme.md

@@ -63,7 +63,7 @@ Local file storage does not require an additional database deployment process, b
 
 #### Deploy NekoImageGallery
 
-1. Clone the project directory to your own PC or server.
+1. Clone the project directory to your own PC or server, then checkout to a specific version tag (like `v1.0.0`).
 2. It is highly recommended to install the dependencies required for this project in a Python venv virtual environment.
    Run the following command:
     ```shell
@@ -80,7 +80,9 @@ Local file storage does not require an additional database deployment process, b
     ```
 5. Modify the project configuration file inside `config/`, you can edit `default.env` directly, but it's recommended to
    create a new file named `local.env` and override the configuration in `default.env`.
-6. Initialize the Qdrant database by running the following command:
+6. ~~Initialize the Qdrant database by running the following command:~~
+   Since NekoImageGallery will now automatically create the collection if not present, this steps is no longer required.
+   However, if you want to explicitly create the collection, you can still run the following command:
     ```shell
     python main.py --init-database
     ```

readme_cn.md

@@ -55,7 +55,7 @@ NekoImageGallery支持两种元数据存储方式：Qdrant数据库存储与本
 
 #### 部署NekoImageGallery
 
-1. 将项目目录clone到你自己的PC或服务器中。
+1. 将项目目录clone到你自己的PC或服务器中，然后按需checkout到特定版本tag（如`v1.0.0`）。
 2. 强烈建议在python venv虚拟环境中安装本项目所需依赖， 运行下面命令：
     ```shell
     python -m venv .venv
@@ -70,7 +70,7 @@ NekoImageGallery支持两种元数据存储方式：Qdrant数据库存储与本
     ```
 5. 按需修改位于`config`目录下的配置文件，您可以直接修改`default.env`，但是建议创建一个名为`local.env`
    的文件，覆盖`default.env`中的配置。
-6. 初始化Qdrant数据库，运行下面命令：
+6. ~~初始化Qdrant数据库，运行下面命令：~~ 由于NekoImageGallery会自动在第一次运行时初始化Qdrant数据库，因此此步骤可省略。但您仍可通过以下命令显式初始化Qdrant数据库：
     ```shell
     python main.py --init-database
     ```

tests/api/integrate_test.py

@@ -38,6 +38,7 @@ async def test_integrate(test_client):
             break
         await asyncio.sleep(1)
 
+    assert resp.json()['index_queue_length'] == 0
     resp = test_client.get('/search/text/hatsune+miku',
                            headers=credentials)
     assert resp.status_code == 200

tests/api/test_upload.py

@@ -0,0 +1,27 @@
+import io
+import random
+
+from .conftest import TEST_ADMIN_TOKEN
+
+
+def test_upload_bad_img_file(test_client):
+    bad_img_file = io.BytesIO(bytearray(random.getrandbits(8) for _ in range(1024 * 1024)))
+    bad_img_file.name = 'bad_image.jpg'
+
+    resp = test_client.post('/admin/upload',
+                            files={'image_file': bad_img_file},
+                            headers={'x-admin-token': TEST_ADMIN_TOKEN},
+                            params={'local': True})
+    print(resp.content)
+    assert resp.status_code == 422
+
+
+def test_upload_unsupported_types(test_client):
+    bad_img_file = io.BytesIO(bytearray(random.getrandbits(8) for _ in range(1024 * 1024)))
+    bad_img_file.name = 'bad_image.tga'
+
+    resp = test_client.post('/admin/upload',
+                            files={'image_file': ('bad_img.tga', bad_img_file, 'image/tga')},
+                            headers={'x-admin-token': TEST_ADMIN_TOKEN},
+                            params={'local': True})
+    assert resp.status_code == 415