Add test for document (infiniflow#3548)

### What problem does this PR solve? Add test for document ### Type of change - [x] New Feature (non-breaking change which adds functionality) Co-authored-by: liuhua <[email protected]>
KevinHuSh · Nov 21, 2024 · c7c8b38 · c7c8b38
1 parent 0ac6dc8
commit c7c8b38
Show file tree

Hide file tree

Showing 21 changed files with 406 additions and 31 deletions.
diff --git a/api/apps/sdk/doc.py b/api/apps/sdk/doc.py
@@ -115,6 +115,7 @@ def upload(dataset_id, tenant_id):
             return get_result(
                 message="No file selected!", code=settings.RetCode.ARGUMENT_ERROR
             )
+    '''
     # total size
     total_size = 0
     for file_obj in file_objs:
@@ -127,6 +128,7 @@ def upload(dataset_id, tenant_id):
             message=f"Total file size exceeds 10MB limit! ({total_size / (1024 * 1024):.2f} MB)",
             code=settings.RetCode.ARGUMENT_ERROR,
         )
+    '''
     e, kb = KnowledgebaseService.get_by_id(dataset_id)
     if not e:
         raise LookupError(f"Can't find the dataset with ID {dataset_id}!")

diff --git a/example/http/simple_example.sh → example/http/dataset_example.sh b/example/http/simple_example.sh → example/http/dataset_example.sh
diff --git a/example/sdk/simple_example.py → example/sdk/dataset_example.py b/example/sdk/simple_example.py → example/sdk/dataset_example.py
diff --git a/sdk/python/test/t_document.py b/sdk/python/test/t_document.py
@@ -1,6 +1,6 @@
-from ragflow_sdk import RAGFlow, DataSet, Document, Chunk
+from ragflow_sdk import RAGFlow
 from common import HOST_ADDRESS
-
+import pytest
 
 def test_upload_document_with_success(get_api_key_fixture):
     API_KEY = get_api_key_fixture
@@ -48,7 +48,6 @@ def test_list_documents_in_dataset_with_success(get_api_key_fixture):
     ds.list_documents(keywords="test", page=1, page_size=12)
 
 
-
 def test_delete_documents_in_dataset_with_success(get_api_key_fixture):
     API_KEY = get_api_key_fixture
     rag = RAGFlow(API_KEY, HOST_ADDRESS)
@@ -59,4 +58,109 @@ def test_delete_documents_in_dataset_with_success(get_api_key_fixture):
     docs = ds.upload_documents(document_infos)
     ds.delete_documents([docs[0].id])
 
+# upload and parse the document with different in different parse method.
+def test_upload_and_parse_pdf_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_pdf_document")
+    with open("test_data/test.pdf","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.pdf","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
 
+def test_upload_and_parse_docx_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_docx_document")
+    with open("test_data/test.docx","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.docx","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
+def test_upload_and_parse_excel_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_excel_document")
+    with open("test_data/test.xlsx","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.xlsx","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
+def test_upload_and_parse_ppt_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_ppt_document")
+    with open("test_data/test.ppt","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.ppt","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
+def test_upload_and_parse_image_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_image_document")
+    with open("test_data/test.jpg","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.jpg","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
+def test_upload_and_parse_txt_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_txt_document")
+    with open("test_data/test.txt","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.txt","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
+def test_upload_and_parse_md_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_md_document")
+    with open("test_data/test.md","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.md","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
+
+def test_upload_and_parse_json_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_json_document")
+    with open("test_data/test.json","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.json","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
+
+@pytest.mark.skip(reason="")
+def test_upload_and_parse_eml_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_eml_document")
+    with open("test_data/test.eml","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.eml","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
+
+def test_upload_and_parse_html_documents_with_general_parse_method(get_api_key_fixture):
+    API_KEY = get_api_key_fixture
+    rag = RAGFlow(API_KEY, HOST_ADDRESS)
+    ds = rag.create_dataset(name="test_html_document")
+    with open("test_data/test.html","rb") as file:
+        blob=file.read()
+    document_infos = [{"displayed_name": "test.html","blob": blob}]
+    docs=ds.upload_documents(document_infos)
+    doc = docs[0]
+    ds.async_parse_documents([doc.id])
diff --git a/sdk/python/test/test_data/.txt b/sdk/python/test/test_data/.txt
diff --git a/sdk/python/test/test_data/empty.txt b/sdk/python/test/test_data/empty.txt
diff --git a/sdk/python/test/test_data/lol.txt b/sdk/python/test/test_data/lol.txt
diff --git a/sdk/python/test/test_data/story.txt b/sdk/python/test/test_data/story.txt
diff --git a/sdk/python/test/test_data/test.docx b/sdk/python/test/test_data/test.docx