57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
#!/usr/bin/env python3
|
|
import json
|
|
import sys
|
|
import time
|
|
from urllib import request, error
|
|
|
|
|
|
def build_chunk(file_url: str, chunk_index: int, chunk_total: int, total_lines: int):
|
|
return {
|
|
"fileUrl": file_url,
|
|
"originalFilename": "demo.xlsx",
|
|
"relativePath": None,
|
|
"mainSheetName": "Sheet1",
|
|
"chunkIndex": chunk_index,
|
|
"chunkTotal": chunk_total,
|
|
"totalLines": total_lines,
|
|
"keptRows": [f"brand-{chunk_index}"],
|
|
"invalidBrands": [],
|
|
"queryFailedBrands": []
|
|
}
|
|
|
|
|
|
def post_json(url: str, payload: dict):
|
|
data = json.dumps(payload).encode("utf-8")
|
|
req = request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
|
|
with request.urlopen(req, timeout=60) as resp:
|
|
return resp.status, resp.read().decode("utf-8", errors="ignore")
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) < 4:
|
|
print("usage: python brand_chunk_load_test.py <base_url> <task_id> <chunk_total> [file_url]")
|
|
sys.exit(1)
|
|
base_url = sys.argv[1].rstrip("/")
|
|
task_id = int(sys.argv[2])
|
|
chunk_total = int(sys.argv[3])
|
|
file_url = sys.argv[4] if len(sys.argv) > 4 else "https://example.com/demo.xlsx"
|
|
|
|
submit_url = f"{base_url}/api/brand/tasks/{task_id}/result"
|
|
started = time.time()
|
|
for i in range(1, chunk_total + 1):
|
|
payload = {
|
|
"strategy": "Terms",
|
|
"files": [build_chunk(file_url, i, chunk_total, chunk_total)]
|
|
}
|
|
status, body = post_json(submit_url, payload)
|
|
if i == 1 or i == chunk_total or i % 100 == 0:
|
|
print(f"chunk {i}/{chunk_total} status={status} elapsed={time.time() - started:.2f}s")
|
|
if status < 200 or status >= 300:
|
|
print(body)
|
|
sys.exit(2)
|
|
print(f"done {chunk_total} chunks in {time.time() - started:.2f}s")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|