From 215bc25f5c848324e540c983b08306db9de2a313 Mon Sep 17 00:00:00 2001
From: zqt <1178747941@qq.com>
Date: Tue, 29 Aug 2023 15:01:26 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E7=9F=A5=E8=AF=86=E5=BA=93?=
 =?UTF-8?q?=E4=BD=BF=E7=94=A8=E7=9A=84=E5=90=91=E9=87=8F=E5=BA=93=E6=B5=8B?=
 =?UTF-8?q?=E8=AF=95=E7=94=A8=E4=BE=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/kb_vector_db/__init__.py       |  0
 tests/kb_vector_db/test_faiss_kb.py  | 35 ++++++++++++++++++++++++++++
 tests/kb_vector_db/test_milvus_db.py | 31 ++++++++++++++++++++++++
 tests/kb_vector_db/test_pg_db.py     | 31 ++++++++++++++++++++++++
 4 files changed, 97 insertions(+)
 create mode 100644 tests/kb_vector_db/__init__.py
 create mode 100644 tests/kb_vector_db/test_faiss_kb.py
 create mode 100644 tests/kb_vector_db/test_milvus_db.py
 create mode 100644 tests/kb_vector_db/test_pg_db.py

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)
+