You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
53 lines
1.4 KiB
53 lines
1.4 KiB
5 months ago
|
from unittest.mock import Mock, patch
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from swarms.memory.qdrant import Qdrant
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_qdrant_client():
|
||
|
|
with patch("swarms.memory.Qdrant") as MockQdrantClient:
|
||
|
|
yield MockQdrantClient()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_sentence_transformer():
|
||
|
|
with patch(
|
||
|
|
"sentence_transformers.SentenceTransformer"
|
||
|
|
) as MockSentenceTransformer:
|
||
|
|
yield MockSentenceTransformer()
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def qdrant_client(mock_qdrant_client, mock_sentence_transformer):
|
||
|
|
client = Qdrant(api_key="your_api_key", host="your_host")
|
||
|
|
yield client
|
||
|
|
|
||
|
|
|
||
|
|
def test_qdrant_init(qdrant_client, mock_qdrant_client):
|
||
|
|
assert qdrant_client.client is not None
|
||
|
|
|
||
|
|
|
||
|
|
def test_load_embedding_model(qdrant_client, mock_sentence_transformer):
|
||
|
|
qdrant_client._load_embedding_model("model_name")
|
||
|
|
mock_sentence_transformer.assert_called_once_with("model_name")
|
||
|
|
|
||
|
|
|
||
|
|
def test_setup_collection(qdrant_client, mock_qdrant_client):
|
||
|
|
qdrant_client._setup_collection()
|
||
|
|
mock_qdrant_client.get_collection.assert_called_once_with(
|
||
|
|
qdrant_client.collection_name
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_add_vectors(qdrant_client, mock_qdrant_client):
|
||
|
|
mock_doc = Mock(page_content="Sample text")
|
||
|
|
qdrant_client.add_vectors([mock_doc])
|
||
|
|
mock_qdrant_client.upsert.assert_called_once()
|
||
|
|
|
||
|
|
|
||
|
|
def test_search_vectors(qdrant_client, mock_qdrant_client):
|
||
|
|
qdrant_client.search_vectors("test query")
|
||
|
|
mock_qdrant_client.search.assert_called_once()
|