Skip to content

Commit

Permalink
test: refine restful testcases trace (milvus-io#34066)
Browse files Browse the repository at this point in the history 
Signed-off-by: zhuwenxing <wenxing.zhu@zilliz.com>
  • Loading branch information
@zhuwenxing
zhuwenxing authored Jun 24, 2024
1 parent b923728 commit 4c6f6c5
Show file tree
Hide file tree
Showing 3 changed files with 412 additions and 41 deletions.
152 changes: 113 additions & 39 deletions tests/restful_client_v2/api/milvus.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,17 +8,75 @@
from minio.commonconfig import CopySource
from tenacity import retry, retry_if_exception_type, stop_after_attempt
from requests.exceptions import ConnectionError


def logger_request_response(response, url, tt, headers, data, str_data, str_response, method):
if len(data) > 2000:
data = data[:1000] + "..." + data[-1000:]
import urllib.parse

ENABLE_LOG_SAVE = False


def simplify_list(lst):
if len(lst) > 20:
return [lst[0], '...', lst[-1]]
return lst


def simplify_dict(d):
if d is None:
d = {}
if len(d) > 20:
keys = list(d.keys())
d = {keys[0]: d[keys[0]], '...': '...', keys[-1]: d[keys[-1]]}
simplified = {}
for k, v in d.items():
if isinstance(v, list):
simplified[k] = simplify_list([simplify_dict(item) if isinstance(item, dict) else simplify_list(
item) if isinstance(item, list) else item for item in v])
elif isinstance(v, dict):
simplified[k] = simplify_dict(v)
else:
simplified[k] = v
return simplified


def build_curl_command(method, url, headers, data=None, params=None):
if isinstance(params, dict):
query_string = urllib.parse.urlencode(params)
url = f"{url}?{query_string}"
curl_cmd = [f"curl -X {method} '{url}'"]

for key, value in headers.items():
curl_cmd.append(f" -H '{key}: {value}'")

if data:
# process_and_simplify(data)
data = json.dumps(data, indent=4)
curl_cmd.append(f" -d '{data}'")

return " \\\n".join(curl_cmd)


def logger_request_response(response, url, tt, headers, data, str_data, str_response, method, params=None):
# save data to jsonl file

data_dict = json.loads(data) if data else {}
data_dict_simple = simplify_dict(data_dict)
if ENABLE_LOG_SAVE:
with open('request_response.jsonl', 'a') as f:
f.write(json.dumps({
"method": method,
"url": url,
"headers": headers,
"params": params,
"data": data_dict_simple,
"response": response.json()
}) + "\n")
data = json.dumps(data_dict_simple, indent=4)
try:
if response.status_code == 200:
if ('code' in response.json() and response.json()["code"] == 0) or (
'Code' in response.json() and response.json()["Code"] == 0):
logger.debug(
f"\nmethod: {method}, \nurl: {url}, \ncost time: {tt}, \nheader: {headers}, \npayload: {str_data}, \nresponse: {str_response}")
f"\nmethod: {method}, \nurl: {url}, \ncost time: {tt}, \nheader: {headers}, \npayload: {data}, \nresponse: {str_response}")

else:
logger.debug(
f"\nmethod: {method}, \nurl: {url}, \ncost time: {tt}, \nheader: {headers}, \npayload: {data}, \nresponse: {response.text}")
Expand All @@ -30,21 +88,31 @@ def logger_request_response(response, url, tt, headers, data, str_data, str_resp
f"method: \nmethod: {method}, \nurl: {url}, \ncost time: {tt}, \nheader: {headers}, \npayload: {data}, \nresponse: {response.text}, \nerror: {e}")


class Requests:
class Requests():
uuid = str(uuid.uuid1())
api_key = None

def __init__(self, url=None, api_key=None):
self.url = url
self.api_key = api_key
if self.uuid is None:
self.uuid = str(uuid.uuid1())
self.headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'RequestId': self.uuid
}

def update_headers(self):
@classmethod
def update_uuid(cls, _uuid):
cls.uuid = _uuid

@classmethod
def update_headers(cls):
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'Authorization': f'Bearer {cls.api_key}',
'RequestId': cls.uuid
}
return headers

Expand All @@ -59,7 +127,7 @@ def post(self, url, headers=None, data=None, params=None):
response = requests.post(url, headers=headers, data=data, params=params)
tt = time.time() - t0
str_response = response.text[:200] + '...' + response.text[-200:] if len(response.text) > 400 else response.text
logger_request_response(response, url, tt, headers, data, str_data, str_response, "post")
logger_request_response(response, url, tt, headers, data, str_data, str_response, "post", params=params)
return response

@retry(retry=retry_if_exception_type(ConnectionError), stop=stop_after_attempt(3))
Expand All @@ -74,7 +142,7 @@ def get(self, url, headers=None, params=None, data=None):
response = requests.get(url, headers=headers, params=params, data=data)
tt = time.time() - t0
str_response = response.text[:200] + '...' + response.text[-200:] if len(response.text) > 400 else response.text
logger_request_response(response, url, tt, headers, data, str_data, str_response, "get")
logger_request_response(response, url, tt, headers, data, str_data, str_response, "get", params=params)
return response

@retry(retry=retry_if_exception_type(ConnectionError), stop=stop_after_attempt(3))
Expand Down Expand Up @@ -111,12 +179,13 @@ def __init__(self, endpoint, token):
self.db_name = None
self.headers = self.update_headers()

def update_headers(self):
@classmethod
def update_headers(cls):
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'Authorization': f'Bearer {cls.api_key}',
'Accept-Type-Allow-Int64': "true",
'RequestId': str(uuid.uuid1())
'RequestId': cls.uuid
}
return headers

Expand Down Expand Up @@ -195,8 +264,6 @@ def vector_hybrid_search(self, payload, db_name="default", timeout=10):

return response.json()



def vector_query(self, payload, db_name="default", timeout=5):
time.sleep(1)
url = f'{self.endpoint}/v2/vectordb/entities/query'
Expand Down Expand Up @@ -269,13 +336,14 @@ def __init__(self, endpoint, token):
self.db_name = None
self.headers = self.update_headers()

def update_headers(self, headers=None):
@classmethod
def update_headers(cls, headers=None):
if headers is not None:
return headers
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'Authorization': f'Bearer {cls.api_key}',
'RequestId': cls.uuid
}
return headers

Expand Down Expand Up @@ -415,11 +483,12 @@ def __init__(self, endpoint, token):
self.db_name = None
self.headers = self.update_headers()

def update_headers(self):
@classmethod
def update_headers(cls):
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'Authorization': f'Bearer {cls.api_key}',
'RequestId': cls.uuid
}
return headers

Expand Down Expand Up @@ -530,11 +599,12 @@ def __init__(self, endpoint, token):
self.db_name = None
self.headers = self.update_headers()

def update_headers(self):
@classmethod
def update_headers(cls):
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'Authorization': f'Bearer {cls.api_key}',
'RequestId': cls.uuid
}
return headers

Expand Down Expand Up @@ -594,11 +664,12 @@ def __init__(self, endpoint, token):
self.headers = self.update_headers()
self.role_names = []

def update_headers(self):
@classmethod
def update_headers(cls):
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'Authorization': f'Bearer {cls.api_key}',
'RequestId': cls.uuid
}
return headers

Expand Down Expand Up @@ -653,11 +724,12 @@ def __init__(self, endpoint, token):
self.db_name = None
self.headers = self.update_headers()

def update_headers(self):
@classmethod
def update_headers(cls):
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'Authorization': f'Bearer {cls.api_key}',
'RequestId': cls.uuid
}
return headers

Expand Down Expand Up @@ -714,11 +786,12 @@ def __init__(self, endpoint, token):
self.db_name = None
self.headers = self.update_headers()

def update_headers(self):
@classmethod
def update_headers(cls):
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'Authorization': f'Bearer {cls.api_key}',
'RequestId': cls.uuid
}
return headers

Expand Down Expand Up @@ -765,11 +838,12 @@ def __init__(self, endpoint, token):
self.db_name = None
self.headers = self.update_headers()

def update_headers(self):
@classmethod
def update_headers(cls):
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {self.api_key}',
'RequestId': str(uuid.uuid1())
'Authorization': f'Bearer {cls.api_key}',
'RequestId': cls.uuid
}
return headers

Expand Down
20 changes: 18 additions & 2 deletions tests/restful_client_v2/base/testbase.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,11 @@
import sys
import pytest
import time
import uuid
from pymilvus import connections, db
from utils.util_log import test_log as logger
from api.milvus import (VectorClient, CollectionClient, PartitionClient, IndexClient, AliasClient,
UserClient, RoleClient, ImportJobClient, StorageClient)
UserClient, RoleClient, ImportJobClient, StorageClient, Requests)
from utils.utils import get_data_by_payload


Expand Down Expand Up @@ -35,7 +36,7 @@ class Base:


class TestBase(Base):

req = None
def teardown_method(self):
self.collection_client.api_key = self.api_key
all_collections = self.collection_client.collection_list()['data']
Expand All @@ -49,19 +50,34 @@ def teardown_method(self):
except Exception as e:
logger.error(e)

# def setup_method(self):
# self.req = Requests()
# self.req.uuid = str(uuid.uuid1())

@pytest.fixture(scope="function", autouse=True)
def init_client(self, endpoint, token, minio_host, bucket_name, root_path):
_uuid = str(uuid.uuid1())
self.req = Requests()
self.req.update_uuid(_uuid)
self.endpoint = f"{endpoint}"
self.api_key = f"{token}"
self.invalid_api_key = "invalid_token"
self.vector_client = VectorClient(self.endpoint, self.api_key)
self.vector_client.update_uuid(_uuid)
self.collection_client = CollectionClient(self.endpoint, self.api_key)
self.collection_client.update_uuid(_uuid)
self.partition_client = PartitionClient(self.endpoint, self.api_key)
self.partition_client.update_uuid(_uuid)
self.index_client = IndexClient(self.endpoint, self.api_key)
self.index_client.update_uuid(_uuid)
self.alias_client = AliasClient(self.endpoint, self.api_key)
self.alias_client.update_uuid(_uuid)
self.user_client = UserClient(self.endpoint, self.api_key)
self.user_client.update_uuid(_uuid)
self.role_client = RoleClient(self.endpoint, self.api_key)
self.role_client.update_uuid(_uuid)
self.import_job_client = ImportJobClient(self.endpoint, self.api_key)
self.import_job_client.update_uuid(_uuid)
self.storage_client = StorageClient(f"{minio_host}:9000", "minioadmin", "minioadmin", bucket_name, root_path)
if token is None:
self.vector_client.api_key = None
Expand Down
Loading

0 comments on commit 4c6f6c5

Please sign in to comment.