diff --git a/tests/kb_vector_db/__init__.py b/tests/kb_vector_db/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/kb_vector_db/test_faiss_kb.py b/tests/kb_vector_db/test_faiss_kb.py new file mode 100644 index 0000000..9c329c8 --- /dev/null +++ b/tests/kb_vector_db/test_faiss_kb.py @@ -0,0 +1,35 @@ +from server.knowledge_base.kb_service.faiss_kb_service import FaissKBService +from server.knowledge_base.migrate import create_tables +from server.knowledge_base.utils import KnowledgeFile + +kbService = FaissKBService("test") +test_kb_name = "test" +test_file_name = "README.md" +testKnowledgeFile = KnowledgeFile(test_file_name, test_kb_name) +search_content = "如何启动api服务" + + +def test_init(): + create_tables() + + +def test_create_db(): + assert kbService.create_kb() + + +def test_add_doc(): + assert kbService.add_doc(testKnowledgeFile) + + +def test_search_db(): + result = kbService.search_docs(search_content) + assert len(result) > 0 +def test_delete_doc(): + assert kbService.delete_doc(testKnowledgeFile) + + + + + +def test_delete_db(): + assert kbService.drop_kb() diff --git a/tests/kb_vector_db/test_milvus_db.py b/tests/kb_vector_db/test_milvus_db.py new file mode 100644 index 0000000..ed72380 --- /dev/null +++ b/tests/kb_vector_db/test_milvus_db.py @@ -0,0 +1,31 @@ +from server.knowledge_base.kb_service.faiss_kb_service import FaissKBService +from server.knowledge_base.kb_service.milvus_kb_service import MilvusKBService +from server.knowledge_base.kb_service.pg_kb_service import PGKBService +from server.knowledge_base.migrate import create_tables +from server.knowledge_base.utils import KnowledgeFile + +kbService = MilvusKBService("test") + +test_kb_name = "test" +test_file_name = "README.md" +testKnowledgeFile = KnowledgeFile(test_file_name, test_kb_name) +search_content = "如何启动api服务" + +def test_init(): + create_tables() + + +def test_create_db(): + assert kbService.create_kb() + + +def test_add_doc(): + assert kbService.add_doc(testKnowledgeFile) + + +def test_search_db(): + result = kbService.search_docs(search_content) + assert len(result) > 0 +def test_delete_doc(): + assert kbService.delete_doc(testKnowledgeFile) + diff --git a/tests/kb_vector_db/test_pg_db.py b/tests/kb_vector_db/test_pg_db.py new file mode 100644 index 0000000..12448d0 --- /dev/null +++ b/tests/kb_vector_db/test_pg_db.py @@ -0,0 +1,31 @@ +from server.knowledge_base.kb_service.faiss_kb_service import FaissKBService +from server.knowledge_base.kb_service.pg_kb_service import PGKBService +from server.knowledge_base.migrate import create_tables +from server.knowledge_base.utils import KnowledgeFile + +kbService = PGKBService("test") + +test_kb_name = "test" +test_file_name = "README.md" +testKnowledgeFile = KnowledgeFile(test_file_name, test_kb_name) +search_content = "如何启动api服务" + + +def test_init(): + create_tables() + + +def test_create_db(): + assert kbService.create_kb() + + +def test_add_doc(): + assert kbService.add_doc(testKnowledgeFile) + + +def test_search_db(): + result = kbService.search_docs(search_content) + assert len(result) > 0 +def test_delete_doc(): + assert kbService.delete_doc(testKnowledgeFile) +