From c7fe75fd1e8008955365030246320b0fe013d283 Mon Sep 17 00:00:00 2001 From: yangxuan Date: Mon, 6 Jan 2025 16:02:31 +0800 Subject: [PATCH] fix ut Signed-off-by: yangxuan --- pymilvus/__init__.py | 2 - tests/test_create_collection.py | 246 -------------------------------- 2 files changed, 248 deletions(-) delete mode 100644 tests/test_create_collection.py diff --git a/pymilvus/__init__.py b/pymilvus/__init__.py index f73aee203..6eb751163 100644 --- a/pymilvus/__init__.py +++ b/pymilvus/__init__.py @@ -14,7 +14,6 @@ from .client.abstract import AnnSearchRequest, Hit, Hits, RRFRanker, SearchResult, WeightedRanker from .client.asynch import SearchFuture from .client.prepare import Prepare -from .client.stub import Milvus from .client.types import ( BulkInsertState, DataType, @@ -89,7 +88,6 @@ "Hits", "Index", "IndexType", - "Milvus", "MilvusClient", "MilvusException", "MilvusUnavailableException", diff --git a/tests/test_create_collection.py b/tests/test_create_collection.py deleted file mode 100644 index 1b6b23425..000000000 --- a/tests/test_create_collection.py +++ /dev/null @@ -1,246 +0,0 @@ -import grpc -import grpc_testing -import pytest -import random -from pymilvus.exceptions import ErrorCode - -from pymilvus.grpc_gen import milvus_pb2, schema_pb2, common_pb2 -from pymilvus import Milvus, DataType - - -class Fields: - class NormalizedField: - def __init__(self, **kwargs): - self.name = kwargs.get("name", None) - self.is_primary_key = kwargs.get("is_primary_key", False) - self.data_type = kwargs.get("data_type", None) - self.type_params = kwargs.get("type_params", dict()) - self.autoID = kwargs.get("autoID", False) - - def __eq__(self, other): - if isinstance(other, Fields.NormalizedField): - return ( - self.name == other.name - and self.is_primary_key == other.is_primary_key - and self.data_type == other.data_type - and self.type_params == other.type_params - and self.autoID == other.autoID - ) - return False - - def __repr__(self): - dump = f"(name: {self.name}" - dump += f", id_primary_key:{self.is_primary_key}" - dump += f", data_type:{self.data_type}" - dump += f", type_params:{self.type_params}" - dump += f", autoID:{self.autoID})" - return dump - - @classmethod - def equal(cls, grpc_fields, dict_fields): - n_grpc_fields = { - field.name: Fields.NormalizedField( - name=field.name, - is_primary_key=field.is_primary_key, - data_type=field.data_type, - type_params={pair.key: pair.value for pair in field.type_params}, - autoID=field.autoID, - ) - for field in grpc_fields - } - n_dict_fields = { - field["name"]: Fields.NormalizedField( - name=field["name"], - is_primary_key=field.get("is_primary", False), - data_type=field["type"], - type_params=field.get("params", dict()), - autoID=field.get("auto_id", False), - ) - for field in dict_fields - } - return n_grpc_fields == n_dict_fields - - -class TestCreateCollection: - @pytest.fixture(scope="function") - def collection_name(self): - return f"test_collection_{random.randint(100000, 999999)}" - - def setup_method(self) -> None: - self._real_time = grpc_testing.strict_real_time() - self._real_time_channel = grpc_testing.channel( - milvus_pb2.DESCRIPTOR.services_by_name.values(), self._real_time - ) - self._servicer = milvus_pb2.DESCRIPTOR.services_by_name["MilvusService"] - self._milvus = Milvus(channel=self._real_time_channel) - - def teardown_method(self) -> None: - pass - - def test_create_collection(self, collection_name): - id_field = { - "name": "my_id", - "type": DataType.INT64, - "auto_id": True, - "is_primary": True, - } - vector_field = { - "name": "embedding", - "type": DataType.FLOAT_VECTOR, - "metric_type": "L2", - "params": {"dim": "4"}, - } - fields = {"fields": [id_field, vector_field], "enable_dynamic_field": True} - future = self._milvus.create_collection( - collection_name=collection_name, fields=fields, _async=True - ) - - invocation_metadata, request, rpc = self._real_time_channel.take_unary_unary( - self._servicer.methods_by_name["CreateCollection"] - ) - rpc.send_initial_metadata(()) - rpc.terminate( - common_pb2.Status( - code=ErrorCode.SUCCESS, error_code=common_pb2.Success, reason="success" - ), - (), - grpc.StatusCode.OK, - "", - ) - - request_schema = schema_pb2.CollectionSchema() - request_schema.ParseFromString(request.schema) - - assert request.collection_name == collection_name - assert Fields.equal(request_schema.fields, fields["fields"]) - assert request_schema.enable_dynamic_field == fields["enable_dynamic_field"] - - return_value = future.result() - assert return_value.code == 0 - assert return_value.reason == "success" - - def test_create_fp16_collection(self, collection_name): - id_field = { - "name": "my_id", - "type": DataType.INT64, - "auto_id": True, - "is_primary": True, - } - vector_field = { - "name": "embedding", - "type": DataType.FLOAT16_VECTOR, - "metric_type": "L2", - "params": {"dim": "4"}, - } - fields = {"fields": [id_field, vector_field], "enable_dynamic_field": True} - future = self._milvus.create_collection( - collection_name=collection_name, fields=fields, _async=True - ) - - invocation_metadata, request, rpc = self._real_time_channel.take_unary_unary( - self._servicer.methods_by_name["CreateCollection"] - ) - rpc.send_initial_metadata(()) - rpc.terminate( - common_pb2.Status( - code=ErrorCode.SUCCESS, error_code=common_pb2.Success, reason="success" - ), - (), - grpc.StatusCode.OK, - "", - ) - - request_schema = schema_pb2.CollectionSchema() - request_schema.ParseFromString(request.schema) - - assert request.collection_name == collection_name - assert Fields.equal(request_schema.fields, fields["fields"]) - assert request_schema.enable_dynamic_field == fields["enable_dynamic_field"] - - return_value = future.result() - assert return_value.code == 0 - assert return_value.reason == "success" - - def test_create_bf16_collection(self, collection_name): - id_field = { - "name": "my_id", - "type": DataType.INT64, - "auto_id": True, - "is_primary": True, - } - vector_field = { - "name": "embedding", - "type": DataType.BFLOAT16_VECTOR, - "metric_type": "L2", - "params": {"dim": "4"}, - } - fields = {"fields": [id_field, vector_field], "enable_dynamic_field": True} - future = self._milvus.create_collection( - collection_name=collection_name, fields=fields, _async=True - ) - - invocation_metadata, request, rpc = self._real_time_channel.take_unary_unary( - self._servicer.methods_by_name["CreateCollection"] - ) - rpc.send_initial_metadata(()) - rpc.terminate( - common_pb2.Status( - code=ErrorCode.SUCCESS, error_code=common_pb2.Success, reason="success" - ), - (), - grpc.StatusCode.OK, - "", - ) - - request_schema = schema_pb2.CollectionSchema() - request_schema.ParseFromString(request.schema) - - assert request.collection_name == collection_name - assert Fields.equal(request_schema.fields, fields["fields"]) - assert request_schema.enable_dynamic_field == fields["enable_dynamic_field"] - - return_value = future.result() - assert return_value.code == 0 - assert return_value.reason == "success" - - def test_create_clustering_key_collection(self, collection_name): - id_field = { - "name": "my_id", - "type": DataType.INT64, - "auto_id": True, - "is_primary": True, - "is_clustering_key": True, - } - vector_field = { - "name": "embedding", - "type": DataType.FLOAT_VECTOR, - "metric_type": "L2", - "params": {"dim": "4"}, - } - fields = {"fields": [id_field, vector_field]} - future = self._milvus.create_collection( - collection_name=collection_name, fields=fields, _async=True - ) - - invocation_metadata, request, rpc = self._real_time_channel.take_unary_unary( - self._servicer.methods_by_name["CreateCollection"] - ) - rpc.send_initial_metadata(()) - rpc.terminate( - common_pb2.Status( - code=ErrorCode.SUCCESS, error_code=common_pb2.Success, reason="success" - ), - (), - grpc.StatusCode.OK, - "", - ) - - request_schema = schema_pb2.CollectionSchema() - request_schema.ParseFromString(request.schema) - - assert request.collection_name == collection_name - assert Fields.equal(request_schema.fields, fields["fields"]) - - return_value = future.result() - assert return_value.code == 0 - assert return_value.reason == "success"