qdrant/tests/openapi/test_timeout.py

from random import random
import subprocess
import pytest

from .helpers.helpers import request_with_validation
from .helpers.collection_setup import drop_collection

num_vectors = 200000
dims = 1024

@pytest.skip("Too big for CI, requires bfb installed", allow_module_level=True)
@pytest.fixture(autouse=True, scope='module')
def setup(collection_name):
    drop_collection(collection_name)
    
    response = request_with_validation(
        api='/collections/{collection_name}',
        method="PUT",
        path_params={'collection_name': collection_name},
        body={
            "vectors": {
                "size": dims,
                "distance": "Euclid",
                "on_disk": True,
            },
            "optimizers_config": {
                "indexing_threshold": 0,
            }
        }
    )
    assert response.ok
    
    subprocess.run(['bfb', '--collection-name', collection_name, '--skip-create', '--dim', str(dims), '--num-vectors', str(num_vectors), '--keywords', '300'])
    
    yield
    drop_collection(collection_name=collection_name)

def test_search_timeout(collection_name):
    response = request_with_validation(
        api='/collections/{collection_name}/points/query',
        method="POST",
        path_params={'collection_name': collection_name},
        query_params={'timeout': 1},
        body={
            "query": [random() for _ in range(dims)],
            "limit": 100000,
            "filter": {
                "must": [
                    {
                        "key": "a",
                        "match": {
                        "text": "keyword_1"
                        }      
                    }
                ]
            },
            "params": { "exact": False }
        }
    )
    
    assert not response.ok
    assert response.status_code == 500
    assert response.json()['status']['error'].__contains__("Timeout error: Operation 'Search' timed out after 1 seconds")

def test_search_batch_timeout(collection_name):
    response = request_with_validation(
        api='/collections/{collection_name}/points/query/batch',
        method="POST",
        path_params={'collection_name': collection_name},
        query_params={'timeout': 1},
        body={
            "searches": [
                {
                    "query": [0.5] * dims,
                    "limit": 100000,
                    "filter": {
                        "must": [
                            {
                                "key": "a",
                                "match": {
                                "text": "keyword_1"
                                }      
                            }
                        ]
                    },
                },
                {
                    "query": [0.6] * dims,
                    "limit": 10000,
                },
            ]
        }
    )
    
    assert not response.ok
    assert response.status_code == 500
    assert response.json()['status']['error'].__contains__("Timeout error: Operation 'Search' timed out after 1 seconds")
    
    
def test_search_groups_timeout(collection_name):
    response = request_with_validation(
        api='/collections/{collection_name}/points/query/groups',
        method="POST",
        path_params={'collection_name': collection_name},
        query_params={'timeout': 1},
        body={
            "query": [0.5] * dims,
            "limit": 100,
            "group_by": "a",
            "filter": {"must": [{"key": "a", "match": {"value": "keyword_1"}}]},
            "group_size": 3,
        }
    )
    
    assert not response.ok
    assert response.status_code == 408
    assert response.json()['status']['error'].__contains__("Timeout error: Operation 'GroupBy' timed out after 1 seconds")
    

def test_recommend_timeout(collection_name):
    response = request_with_validation(
        api='/collections/{collection_name}/points/query',
        method="POST",
        path_params={'collection_name': collection_name},
        query_params={'timeout': 1},
        body={
            "query": {
                "recommend": {
                    "positive": [111,222,333,444,555],
                    "negative": [666,777,888,999,1010],
                    "strategy": "best_score",
                }
            },
            "limit": 10000,
        }
    )
    
    assert not response.ok
    assert response.status_code == 500
    assert response.json()['status']['error'].__contains__("Timeout error: Operation 'Search' timed out after 1 seconds")
    
    
def test_recommend_batch_timeout(collection_name):
    response = request_with_validation(
        api='/collections/{collection_name}/points/query/batch',
        method="POST",
        path_params={'collection_name': collection_name},
        query_params={'timeout': 1},
        body={
            "searches": [
                {
                    "query": {
                        "recommend": {
                            "positive": [111,222,333,444,555],
                            "negative": [666,777,888,999,1010],
                            "strategy": "best_score",
                        }
                    },
                    "limit": 10000,
                },
                {
                    "query": {
                        "recommend": {
                            "positive": [666,777,888,999,1010],
                            "negative": [111,222,333,444,555],
                            "strategy": "best_score",
                        }
                    },
                    "limit": 10000,
                },
            ]
        }
    )
    
    assert not response.ok
    assert response.status_code == 500
    assert response.json()['status']['error'].__contains__("Timeout error: Operation 'Search' timed out after 1 seconds")
    

def test_recommend_groups_timeout(collection_name):
    response = request_with_validation(
        api='/collections/{collection_name}/points/query/groups',
        method="POST",
        path_params={'collection_name': collection_name},
        query_params={'timeout': 1},
        body={
            "query": {
                "recommend": {
                    "positive": [111,555],
                    "negative": [666,1010],
                }
            },
            "limit": 100,
            "filter": {"must": [{"key": "a", "match": {"value": "keyword_1"}}]},
            "group_by": "a",
            "group_size": 3,
        }
    )
    
    assert not response.ok
    assert response.status_code == 408
    assert response.json()['status']['error'].__contains__("Timeout error: Operation 'GroupBy' timed out after 1 seconds")
    
    
def test_discover_timeout(collection_name):
    response = request_with_validation(
        api='/collections/{collection_name}/points/query',
        method="POST",
        path_params={'collection_name': collection_name},
        query_params={'timeout': 1},
        body={
            "query": {
                "discover": {
                    "target": 111,
                    "context": [
                        {"positive": 666, "negative": 777},
                        {"positive": 888, "negative": 999},
                        {"positive": 1010, "negative": 1111},
                    ],
                }
            },
            "limit": 100,
        }
    )
    
    assert not response.ok
    assert response.status_code == 500
    assert response.json()['status']['error'].__contains__("Timeout error: Operation 'Search' timed out after 1 seconds")