mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 07:38:29 +00:00
95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
"""Test Elasticsearch knowledge storage functionality."""
|
|
|
|
import os
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from crewai.knowledge.storage.elasticsearch_knowledge_storage import (
|
|
ElasticsearchKnowledgeStorage,
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.environ.get("RUN_ELASTICSEARCH_TESTS") != "true",
|
|
reason="Elasticsearch tests require RUN_ELASTICSEARCH_TESTS=true"
|
|
)
|
|
class TestElasticsearchKnowledgeStorage(unittest.TestCase):
|
|
"""Test Elasticsearch knowledge storage functionality."""
|
|
|
|
def setUp(self):
|
|
"""Set up test fixtures."""
|
|
self.es_mock = MagicMock()
|
|
self.es_mock.indices.exists.return_value = False
|
|
|
|
self.embedder_mock = MagicMock()
|
|
self.embedder_mock.embed_documents.return_value = [[0.1, 0.2, 0.3]]
|
|
|
|
self.es_patcher = patch(
|
|
"crewai.knowledge.storage.elasticsearch_knowledge_storage.Elasticsearch",
|
|
return_value=self.es_mock
|
|
)
|
|
self.es_class_mock = self.es_patcher.start()
|
|
|
|
self.storage = ElasticsearchKnowledgeStorage(
|
|
embedder_config=self.embedder_mock,
|
|
collection_name="test"
|
|
)
|
|
self.storage.initialize_knowledge_storage()
|
|
|
|
def tearDown(self):
|
|
"""Tear down test fixtures."""
|
|
self.es_patcher.stop()
|
|
|
|
def test_initialization(self):
|
|
"""Test initialization of Elasticsearch knowledge storage."""
|
|
self.es_class_mock.assert_called_once()
|
|
|
|
self.es_mock.indices.create.assert_called_once()
|
|
|
|
def test_save(self):
|
|
"""Test saving to Elasticsearch knowledge storage."""
|
|
self.storage.save(["Test document 1", "Test document 2"], {"source": "test"})
|
|
|
|
self.assertEqual(self.es_mock.index.call_count, 2)
|
|
|
|
self.assertEqual(self.embedder_mock.embed_documents.call_count, 2)
|
|
|
|
def test_search(self):
|
|
"""Test searching in Elasticsearch knowledge storage."""
|
|
self.es_mock.search.return_value = {
|
|
"hits": {
|
|
"hits": [
|
|
{
|
|
"_id": "test_id",
|
|
"_score": 1.5, # Score between 1-2 (Elasticsearch range)
|
|
"_source": {
|
|
"text": "Test document",
|
|
"metadata": {"source": "test"},
|
|
}
|
|
}
|
|
]
|
|
}
|
|
}
|
|
|
|
results = self.storage.search(["test query"])
|
|
|
|
self.es_mock.search.assert_called_once()
|
|
|
|
self.assertEqual(len(results), 1)
|
|
self.assertEqual(results[0]["id"], "test_id")
|
|
self.assertEqual(results[0]["context"], "Test document")
|
|
self.assertEqual(results[0]["metadata"], {"source": "test"})
|
|
self.assertEqual(results[0]["score"], 0.5) # Adjusted to 0-1 range
|
|
|
|
def test_reset(self):
|
|
"""Test resetting Elasticsearch knowledge storage."""
|
|
self.es_mock.indices.exists.return_value = True
|
|
|
|
self.storage.reset()
|
|
|
|
self.es_mock.indices.delete.assert_called_once()
|
|
|
|
self.assertEqual(self.es_mock.indices.create.call_count, 2)
|