legal-doc-masker/magicdoc/test_api.py

93 lines
2.9 KiB
Python

#!/usr/bin/env python3
"""
Test script for MagicDoc API
"""
import requests
import json
import os
def test_health_check(base_url="http://localhost:8002"):
"""Test health check endpoint"""
try:
response = requests.get(f"{base_url}/health")
print(f"Health check status: {response.status_code}")
print(f"Response: {response.json()}")
return response.status_code == 200
except Exception as e:
print(f"Health check failed: {e}")
return False
def test_file_parse(base_url="http://localhost:8002", file_path=None):
"""Test file parse endpoint"""
if not file_path or not os.path.exists(file_path):
print(f"File not found: {file_path}")
return False
try:
with open(file_path, 'rb') as f:
files = {'files': (os.path.basename(file_path), f, 'application/octet-stream')}
data = {
'output_dir': './output',
'lang_list': 'ch',
'backend': 'pipeline',
'parse_method': 'auto',
'formula_enable': True,
'table_enable': True,
'return_md': True,
'return_middle_json': False,
'return_model_output': False,
'return_content_list': False,
'return_images': False,
'start_page_id': 0,
'end_page_id': 99999
}
response = requests.post(f"{base_url}/file_parse", files=files, data=data)
print(f"File parse status: {response.status_code}")
if response.status_code == 200:
result = response.json()
print(f"Success! Converted {len(result.get('markdown', ''))} characters")
print(f"Time cost: {result.get('time_cost', 'N/A')}s")
return True
else:
print(f"Error: {response.text}")
return False
except Exception as e:
print(f"File parse failed: {e}")
return False
def main():
"""Main test function"""
print("Testing MagicDoc API...")
# Test health check
print("\n1. Testing health check...")
if not test_health_check():
print("Health check failed. Make sure the service is running.")
return
# Test file parse (if sample file exists)
print("\n2. Testing file parse...")
sample_files = [
"../sample_doc/20220707_na_decision-2.docx",
"../sample_doc/20220707_na_decision-2.pdf",
"../sample_doc/short_doc.md"
]
for sample_file in sample_files:
if os.path.exists(sample_file):
print(f"Testing with {sample_file}...")
if test_file_parse(file_path=sample_file):
print("File parse test passed!")
break
else:
print(f"Sample file not found: {sample_file}")
print("\nTest completed!")
if __name__ == "__main__":
main()