Sidney233 3 miesięcy temu
rodzic
commit
84a9e5f13e
1 zmienionych plików z 10 dodań i 4 usunięć
  1. 10 4
      tests/unittest/test_e2e.py

+ 10 - 4
tests/unittest/test_e2e.py

@@ -36,7 +36,7 @@ def test_pipeline_with_two_config():
         if doc_path.suffix in pdf_suffixes + image_suffixes:
             doc_path_list.append(doc_path)
 
-    # os.environ["MINERU_MODEL_SOURCE"] = "modelscope"
+    os.environ["MINERU_MODEL_SOURCE"] = "local"
 
     pdf_file_names = []
     pdf_bytes_list = []
@@ -69,7 +69,10 @@ def test_pipeline_with_two_config():
         output_dir,
         parse_method="txt",
     )
-    assert_content("tests/unittest/output/test/txt/test_content_list.json")
+    res_json_path = (
+        Path(__file__).parent / "output" / "test" / "txt" / "test_content_list.json"
+    ).as_posix()
+    assert_content(res_json_path)
     infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = (
         pipeline_doc_analyze(
             pdf_bytes_list,
@@ -102,7 +105,7 @@ def test_vlm_transformers_with_default_config():
         if doc_path.suffix in pdf_suffixes + image_suffixes:
             doc_path_list.append(doc_path)
 
-    # os.environ["MINERU_MODEL_SOURCE"] = "modelscope"
+    os.environ["MINERU_MODEL_SOURCE"] = "modelscope"
 
     pdf_file_names = []
     pdf_bytes_list = []
@@ -155,7 +158,10 @@ def test_vlm_transformers_with_default_config():
         )
 
         logger.info(f"local output dir is {local_md_dir}")
-        assert_content("tests/unittest/output/test/vlm/test_content_list.json")
+        res_json_path = (
+            Path(__file__).parent / "output" / "test" / "txt" / "test_content_list.json"
+        ).as_posix()
+        assert_content(res_json_path)
 
 
 def write_infer_result(