109 lines
4.2 KiB
Python
109 lines
4.2 KiB
Python
import requests
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
# Configuration
|
|
BASE_URL = "http://localhost:3015"
|
|
API_KEY = "jleu1212"
|
|
OCR_PDF_PATH = "ocr.pdf"
|
|
|
|
def test_ocr_upload_with_api_key():
|
|
print("Testing OCR PDF upload with API key authentication...")
|
|
print(f"API Key: {API_KEY}")
|
|
print(f"OCR PDF: {OCR_PDF_PATH}")
|
|
|
|
# Use query parameter for API key as per OpenAPI spec
|
|
params = {"api_key_header_value": API_KEY}
|
|
|
|
# Step 1: Check server status
|
|
print("\n1. Checking server status...")
|
|
try:
|
|
status_response = requests.get(f"{BASE_URL}/", params=params)
|
|
if status_response.status_code == 200:
|
|
print("✓ Server is running and accessible")
|
|
else:
|
|
print(f"⚠ Server status: {status_response.status_code}")
|
|
except Exception as e:
|
|
print(f"✗ Server check error: {e}")
|
|
return False
|
|
|
|
# Step 2: Clear existing documents first
|
|
print("\n2. Clearing existing documents...")
|
|
try:
|
|
clear_response = requests.delete(f"{BASE_URL}/documents", params=params)
|
|
if clear_response.status_code == 200:
|
|
clear_result = clear_response.json()
|
|
print(f"✓ Documents cleared successfully: {clear_result}")
|
|
else:
|
|
print(f"⚠ Clear documents response: {clear_response.status_code} - {clear_response.text}")
|
|
except Exception as e:
|
|
print(f"⚠ Clear documents error: {e}")
|
|
|
|
# Step 3: Upload OCR PDF
|
|
print("\n3. Uploading OCR PDF...")
|
|
|
|
if not Path(OCR_PDF_PATH).exists():
|
|
print(f"✗ OCR PDF file not found: {OCR_PDF_PATH}")
|
|
return False
|
|
|
|
files = {"file": (OCR_PDF_PATH, open(OCR_PDF_PATH, "rb"), "application/pdf")}
|
|
|
|
try:
|
|
upload_response = requests.post(f"{BASE_URL}/documents/upload", files=files, params=params)
|
|
print(f"Upload response status: {upload_response.status_code}")
|
|
print(f"Upload response text: {upload_response.text}")
|
|
|
|
if upload_response.status_code == 200:
|
|
upload_result = upload_response.json()
|
|
print(f"✓ Upload successful! Result: {json.dumps(upload_result, indent=2)}")
|
|
|
|
# Step 4: Wait for processing and check document status
|
|
print("\n4. Waiting for document processing...")
|
|
time.sleep(10) # Give more time for OCR processing
|
|
|
|
try:
|
|
status_response = requests.get(f"{BASE_URL}/documents", params=params)
|
|
if status_response.status_code == 200:
|
|
documents = status_response.json()
|
|
print(f"✓ Documents status: {json.dumps(documents, indent=2)}")
|
|
else:
|
|
print(f"⚠ Documents status check failed: {status_response.status_code}")
|
|
except Exception as e:
|
|
print(f"⚠ Documents status check error: {e}")
|
|
|
|
# Step 5: Test search functionality
|
|
print("\n5. Testing search functionality...")
|
|
search_data = {
|
|
"query": "safety precautions",
|
|
"top_k": 5
|
|
}
|
|
|
|
try:
|
|
search_response = requests.post(f"{BASE_URL}/query", json=search_data, params=params)
|
|
if search_response.status_code == 200:
|
|
search_results = search_response.json()
|
|
print(f"✓ Search successful! Response: {json.dumps(search_results, indent=2)}")
|
|
else:
|
|
print(f"⚠ Search failed: {search_response.status_code} - {search_response.text}")
|
|
except Exception as e:
|
|
print(f"⚠ Search error: {e}")
|
|
|
|
return True
|
|
else:
|
|
print(f"✗ Upload failed: {upload_response.status_code} - {upload_response.text}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"✗ Upload error: {e}")
|
|
return False
|
|
finally:
|
|
if 'files' in locals():
|
|
files["file"][1].close()
|
|
|
|
if __name__ == "__main__":
|
|
success = test_ocr_upload_with_api_key()
|
|
if success:
|
|
print("\n🎉 OCR upload test with API key authentication completed successfully!")
|
|
else:
|
|
print("\n❌ OCR upload test with API key authentication failed!") |