74 lines
2.6 KiB
Python
74 lines
2.6 KiB
Python
import requests
|
|
import time
|
|
import sys
|
|
|
|
SERVER_URL = "http://localhost:3015"
|
|
UPLOAD_ENDPOINT = f"{SERVER_URL}/documents/upload"
|
|
SEARCH_ENDPOINT = f"{SERVER_URL}/search"
|
|
STATUS_ENDPOINT = f"{SERVER_URL}/documents/status"
|
|
API_KEY = "jleu1212"
|
|
|
|
def get_headers(workspace=None):
|
|
headers = {"X-API-Key": API_KEY}
|
|
if workspace:
|
|
headers["X-Workspace"] = workspace
|
|
return headers
|
|
|
|
def test_upload(file_path, workspace=None):
|
|
"""Test upload with optional workspace header"""
|
|
headers = get_headers(workspace)
|
|
with open(file_path, 'rb') as f:
|
|
files = {'file': (file_path.name, f)}
|
|
resp = requests.post(UPLOAD_ENDPOINT, files=files, headers=headers)
|
|
print(f"Upload response: {resp.status_code} {resp.text}")
|
|
return resp.json() if resp.status_code == 200 else None
|
|
|
|
def test_search(query, workspace=None):
|
|
headers = get_headers(workspace)
|
|
resp = requests.post(SEARCH_ENDPOINT, json={"query": query}, headers=headers)
|
|
print(f"Search response: {resp.status_code} {resp.text}")
|
|
return resp.json() if resp.status_code == 200 else None
|
|
|
|
def test_status(workspace=None):
|
|
headers = get_headers(workspace)
|
|
resp = requests.get(STATUS_ENDPOINT, headers=headers)
|
|
print(f"Status response: {resp.status_code} {resp.text}")
|
|
return resp.json() if resp.status_code == 200 else None
|
|
|
|
if __name__ == "__main__":
|
|
from pathlib import Path
|
|
test_file = Path("test/test.docx")
|
|
if not test_file.exists():
|
|
print("Test file not found")
|
|
sys.exit(1)
|
|
|
|
# Test upload to workspace "test1"
|
|
print("=== Upload to workspace test1 ===")
|
|
result = test_upload(test_file, workspace="test1")
|
|
if result:
|
|
track_id = result.get("track_id")
|
|
print(f"Track ID: {track_id}")
|
|
# Wait a bit for processing
|
|
time.sleep(5)
|
|
# Check status
|
|
print("=== Status for workspace test1 ===")
|
|
test_status("test1")
|
|
# Search
|
|
print("=== Search in workspace test1 ===")
|
|
test_search("test", workspace="test1")
|
|
|
|
# Test upload to workspace "test2"
|
|
print("\n=== Upload to workspace test2 ===")
|
|
result = test_upload(test_file, workspace="test2")
|
|
if result:
|
|
track_id = result.get("track_id")
|
|
print(f"Track ID: {track_id}")
|
|
time.sleep(5)
|
|
print("=== Status for workspace test2 ===")
|
|
test_status("test2")
|
|
print("=== Search in workspace test2 ===")
|
|
test_search("test", workspace="test2")
|
|
|
|
# Test search without workspace (should default to default?)
|
|
print("\n=== Search without workspace (should be default?) ===")
|
|
test_search("test") |