diff --git a/tests/restful_client_v2/api/milvus.py b/tests/restful_client_v2/api/milvus.py index 596f9f9b62..c31eb97efa 100644 --- a/tests/restful_client_v2/api/milvus.py +++ b/tests/restful_client_v2/api/milvus.py @@ -92,17 +92,16 @@ def logger_request_response(response, url, tt, headers, data, str_data, str_resp 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()) + if self.__class__.uuid is None: + self.__class__.uuid = str(uuid.uuid1()) self.headers = { 'Content-Type': 'application/json', 'Authorization': f'Bearer {self.api_key}', - 'RequestId': self.uuid, + 'RequestId': self.__class__.uuid, "Request-Timeout": REQUEST_TIMEOUT } @@ -110,12 +109,11 @@ class Requests(): def update_uuid(cls, _uuid): cls.uuid = _uuid - @classmethod - def update_headers(cls): + def update_headers(self): headers = { 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', - 'RequestId': cls.uuid, + 'Authorization': f'Bearer {self.api_key}', + 'RequestId': self.__class__.uuid, "Request-Timeout": REQUEST_TIMEOUT } return headers @@ -183,13 +181,12 @@ class VectorClient(Requests): self.db_name = None self.headers = self.update_headers() - @classmethod - def update_headers(cls): + def update_headers(self): headers = { 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', + 'Authorization': f'Bearer {self.api_key}', 'Accept-Type-Allow-Int64': "true", - 'RequestId': cls.uuid, + 'RequestId': self.__class__.uuid, "Request-Timeout": REQUEST_TIMEOUT } return headers @@ -352,14 +349,13 @@ class CollectionClient(Requests): else: time.sleep(1) - @classmethod - def update_headers(cls, headers=None): + def update_headers(self, headers=None): if headers is not None: return headers headers = { 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', - 'RequestId': cls.uuid, + 'Authorization': f'Bearer {self.api_key}', + 'RequestId': self.__class__.uuid, "Request-Timeout": REQUEST_TIMEOUT } return headers @@ -613,16 +609,6 @@ class PartitionClient(Requests): self.db_name = None self.headers = self.update_headers() - @classmethod - def update_headers(cls): - headers = { - 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', - 'RequestId': cls.uuid, - "Request-Timeout": REQUEST_TIMEOUT - } - return headers - def partition_list(self, db_name="default", collection_name=None): url = f'{self.endpoint}/v2/vectordb/partitions/list' data = { @@ -730,15 +716,6 @@ class UserClient(Requests): self.db_name = None self.headers = self.update_headers() - @classmethod - def update_headers(cls): - headers = { - 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', - 'RequestId': cls.uuid - } - return headers - def user_list(self): url = f'{self.endpoint}/v2/vectordb/users/list' response = self.post(url, headers=self.update_headers()) @@ -795,15 +772,6 @@ class RoleClient(Requests): self.headers = self.update_headers() self.role_names = [] - @classmethod - def update_headers(cls): - headers = { - 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', - 'RequestId': cls.uuid - } - return headers - def role_list(self): url = f'{self.endpoint}/v2/vectordb/roles/list' response = self.post(url, headers=self.update_headers()) @@ -855,16 +823,6 @@ class IndexClient(Requests): self.db_name = None self.headers = self.update_headers() - @classmethod - def update_headers(cls): - headers = { - 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', - 'RequestId': cls.uuid, - "Request-Timeout": REQUEST_TIMEOUT - } - return headers - def index_create(self, payload, db_name="default"): url = f'{self.endpoint}/v2/vectordb/indexes/create' if self.db_name is not None: @@ -948,15 +906,6 @@ class AliasClient(Requests): self.db_name = None self.headers = self.update_headers() - @classmethod - def update_headers(cls): - headers = { - 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', - 'RequestId': cls.uuid - } - return headers - def list_alias(self): url = f'{self.endpoint}/v2/vectordb/aliases/list' response = self.post(url, headers=self.update_headers()) @@ -1000,16 +949,6 @@ class ImportJobClient(Requests): self.db_name = None self.headers = self.update_headers() - @classmethod - def update_headers(cls): - headers = { - 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}', - 'RequestId': cls.uuid, - "Request-Timeout": REQUEST_TIMEOUT - } - return headers - def list_import_jobs(self, payload, db_name="default"): if self.db_name is not None: db_name = self.db_name @@ -1069,14 +1008,6 @@ class DatabaseClient(Requests): self.db_name = None self.db_names = [] # Track created databases - @classmethod - def update_headers(cls): - headers = { - 'Content-Type': 'application/json', - 'Authorization': f'Bearer {cls.api_key}' - } - return headers - def database_create(self, payload): """Create a database""" url = f"{self.endpoint}/v2/vectordb/databases/create" diff --git a/tests/restful_client_v2/base/testbase.py b/tests/restful_client_v2/base/testbase.py index e727f087fd..e685cc0f29 100644 --- a/tests/restful_client_v2/base/testbase.py +++ b/tests/restful_client_v2/base/testbase.py @@ -42,7 +42,8 @@ class TestBase(Base): def teardown_method(self): # Clean up collections - self.collection_client.api_key = self.api_key + if hasattr(self, 'api_key') and self.api_key: + self.collection_client.api_key = self.api_key all_collections = self.collection_client.collection_list()['data'] if self.name in all_collections: logger.info(f"collection {self.name} exist, drop it") @@ -68,7 +69,8 @@ class TestBase(Base): # Clean up databases created by this client - self.database_client.api_key = self.api_key + if hasattr(self, 'api_key') and self.api_key: + self.database_client.api_key = self.api_key for db_name in self.database_client.db_names[:]: # Create a copy of the list to iterate logger.info(f"database {db_name} exist, drop it") try: @@ -84,25 +86,19 @@ class TestBase(Base): self.endpoint = f"{endpoint}" self.api_key = f"{token}" self.invalid_api_key = "invalid_token" + Requests.update_uuid(_uuid) + 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) self.database_client = DatabaseClient(self.endpoint, self.api_key) - self.database_client.update_uuid(_uuid) + if token is None: self.vector_client.api_key = None self.collection_client.api_key = None