SUIME/test.py

93 lines
2.7 KiB
Python

#!/usr/bin/env python3
import socket
import struct
import msgpack
import sys
import os
import time
# Define test data: 5 requests with different pinyin and context
test_requests = [
{"pinyin": "shanghai", "context": ""},
{"pinyin": "beijing", "context": ""},
{"pinyin": "nihao", "context": ""},
{"pinyin": "xiexie", "context": ""},
{"pinyin": "f", "context": "返回汉字的完整拼音和剩余部"}
]
# Determine socket address based on operating system
def get_socket_address():
if sys.platform == "win32":
# Windows: TCP socket
return ("127.0.0.1", 23333), "tcp"
else:
# Unix-like systems: Unix Domain Socket
return "/tmp/su-ime.sock", "unix"
# Connect to server
def connect_to_server(address, socket_type):
if socket_type == "unix":
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
sock.connect(address)
else: # tcp
host, port = address
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect((host, port))
return sock
# Send request and receive response
def send_request(sock, request):
# Serialize request to MessagePack
payload = msgpack.packb(request)
payload_len = len(payload)
# Create header: 4-byte big-endian unsigned integer
header = struct.pack(">I", payload_len)
# Send header and payload
sock.sendall(header + payload)
# Receive response
# Read header
len_buf = sock.recv(4)
if len(len_buf) != 4:
raise ValueError("Failed to read response header")
response_len = struct.unpack(">I", len_buf)[0]
# Read payload
response_payload = b""
while len(response_payload) < response_len:
chunk = sock.recv(response_len - len(response_payload))
if not chunk:
break
response_payload += chunk
# Deserialize response
response = msgpack.unpackb(response_payload, raw=False)
return response
# Main function
def main():
address, socket_type = get_socket_address()
try:
sock = connect_to_server(address, socket_type)
print(f"Connected to server via {socket_type} at {address}")
for i, req in enumerate(test_requests, 1):
print(f"\nSending request {i}: {req}")
try:
start = time.time()
response = send_request(sock, req)
print(f'cost time: {time.time() - start}')
print(f"Response {i}: {response[0][0:10]}")
except Exception as e:
print(f"Error sending request {i}: {e}")
sock.close()
print("\nAll requests sent.")
except Exception as e:
print(f"Connection error: {e}")
sys.exit(1)
if __name__ == "__main__":
main()