|
@@ -0,0 +1,353 @@
|
|
|
+import {
|
|
|
+ AlterAliasReq,
|
|
|
+ CreateAliasReq,
|
|
|
+ CreateCollectionReq,
|
|
|
+ CreateIndexReq,
|
|
|
+ CreatePartitionReq,
|
|
|
+ DescribeCollectionReq,
|
|
|
+ DescribeIndexReq,
|
|
|
+ DropAliasReq,
|
|
|
+ DropCollectionReq,
|
|
|
+ DropIndexReq,
|
|
|
+ DropPartitionReq,
|
|
|
+ FlushReq,
|
|
|
+ GetCollectionStatisticsReq,
|
|
|
+ GetIndexBuildProgressReq,
|
|
|
+ GetIndexStateReq,
|
|
|
+ GetPartitionStatisticsReq,
|
|
|
+ InsertReq,
|
|
|
+ LoadCollectionReq,
|
|
|
+ LoadPartitionsReq,
|
|
|
+ ReleaseLoadCollectionReq,
|
|
|
+ ReleasePartitionsReq,
|
|
|
+ SearchReq,
|
|
|
+ ShowCollectionsReq,
|
|
|
+ ShowPartitionsReq,
|
|
|
+} from '@zilliz/milvus2-sdk-node/dist/milvus/types';
|
|
|
+import { DeleteEntitiesReq } from '@zilliz/milvus2-sdk-node/dist/milvus/types/Data';
|
|
|
+import { QueryDto } from '../../../collections/dto';
|
|
|
+import {
|
|
|
+ CodeEnum,
|
|
|
+ ERR_NO_ADDRESS,
|
|
|
+ ERR_NO_ALIAS,
|
|
|
+ ERR_NO_COLLECTION,
|
|
|
+ ERR_NO_INDEX,
|
|
|
+ ERR_NO_PARAM,
|
|
|
+} from '../../utils/constants';
|
|
|
+import { mockStatusInfo } from '../../utils/mock.util';
|
|
|
+import {
|
|
|
+ mockCollectionNames,
|
|
|
+ mockCollections,
|
|
|
+ mockIndexState,
|
|
|
+ mockLoadedCollections,
|
|
|
+ mockPartition,
|
|
|
+} from '../consts';
|
|
|
+
|
|
|
+const mockMilvusClient = jest.fn().mockImplementation((address: string) => {
|
|
|
+ return {
|
|
|
+ collectionManager: {
|
|
|
+ hasCollection: (param: { collection_name: string }) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (address === '') {
|
|
|
+ throw new Error(ERR_NO_ADDRESS);
|
|
|
+ }
|
|
|
+ return collection_name;
|
|
|
+ },
|
|
|
+ showCollections: (param?: ShowCollectionsReq) => {
|
|
|
+ if (!param) {
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: mockCollectionNames,
|
|
|
+ };
|
|
|
+ }
|
|
|
+ const { collection_names, type } = param;
|
|
|
+ // loaded type
|
|
|
+ if (type === 1) {
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: mockLoadedCollections,
|
|
|
+ };
|
|
|
+ }
|
|
|
+ return collection_names && collection_names.length > 0
|
|
|
+ ? {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: collection_names,
|
|
|
+ }
|
|
|
+ : { status: mockStatusInfo(CodeEnum.error, ERR_NO_PARAM) };
|
|
|
+ },
|
|
|
+ createCollection: (param: CreateCollectionReq) => {
|
|
|
+ const { collection_name, fields } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+ return { ...mockStatusInfo(CodeEnum.success), data: fields };
|
|
|
+ },
|
|
|
+ describeCollection: (param: DescribeCollectionReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION),
|
|
|
+ };
|
|
|
+ }
|
|
|
+ const res =
|
|
|
+ mockCollections.find((c) => c.name === collection_name) || {};
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ ...res,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ dropCollection: (param: DropCollectionReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+ return { ...mockStatusInfo(CodeEnum.success), data: collection_name };
|
|
|
+ },
|
|
|
+ loadCollection: (param: LoadCollectionReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+ return { ...mockStatusInfo(CodeEnum.success), data: collection_name };
|
|
|
+ },
|
|
|
+ releaseCollection: (param: ReleaseLoadCollectionReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+
|
|
|
+ return { ...mockStatusInfo(CodeEnum.success), data: collection_name };
|
|
|
+ },
|
|
|
+ getCollectionStatistics: (param: GetCollectionStatisticsReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+ const data = {
|
|
|
+ name: collection_name,
|
|
|
+ stats: [{ key: 'row_count', value: 7 }],
|
|
|
+ };
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ ...data,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ createAlias: (param: CreateAliasReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ ...mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ alterAlias: (param: AlterAliasReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ ...mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ dropAlias: (param: DropAliasReq) => {
|
|
|
+ const { alias } = param;
|
|
|
+ if (!alias) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_ALIAS);
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ ...mockStatusInfo(CodeEnum.success),
|
|
|
+ data: alias,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ },
|
|
|
+ partitionManager: {
|
|
|
+ createPartition: (param: CreatePartitionReq) => {
|
|
|
+ const { collection_name, partition_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+ return { ...mockStatusInfo(CodeEnum.success), data: partition_name };
|
|
|
+ },
|
|
|
+ dropPartition: (param: DropPartitionReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+ return { ...mockStatusInfo(CodeEnum.success), data: param };
|
|
|
+ },
|
|
|
+ loadPartitions: (param: LoadPartitionsReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+ return { ...mockStatusInfo(CodeEnum.success), data: param };
|
|
|
+ },
|
|
|
+ releasePartitions: (param: ReleasePartitionsReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+ return { ...mockStatusInfo(CodeEnum.success), data: param };
|
|
|
+ },
|
|
|
+ showPartitions: (param: ShowPartitionsReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ ...mockPartition,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ getPartitionStatistics: (param: GetPartitionStatisticsReq) => {
|
|
|
+ const { collection_name, partition_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+
|
|
|
+ const data = {
|
|
|
+ name: partition_name,
|
|
|
+ stats: [{ key: 'row_count', value: 7 }],
|
|
|
+ };
|
|
|
+
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ ...data,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ },
|
|
|
+ indexManager: {
|
|
|
+ getIndexState: (param: GetIndexStateReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+ const data =
|
|
|
+ mockIndexState.find((i) => i.collection_name === collection_name) ||
|
|
|
+ {};
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ ...data,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ createIndex: (param: CreateIndexReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ ...mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ describeIndex: (param: DescribeIndexReq) => {
|
|
|
+ const { collection_name, field_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+ if (!field_name) {
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.indexNoExist, ERR_NO_INDEX),
|
|
|
+ };
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ dropIndex: (param: DropIndexReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ ...mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ getIndexBuildProgress: (param: GetIndexBuildProgressReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ },
|
|
|
+ dataManager: {
|
|
|
+ flush: (data: FlushReq) => ({
|
|
|
+ data,
|
|
|
+ }),
|
|
|
+ getMetric: (data: { request: { metric_type: string } }) => {
|
|
|
+ const {
|
|
|
+ request: { metric_type: type },
|
|
|
+ } = data;
|
|
|
+
|
|
|
+ return {
|
|
|
+ type,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ insert: (param: InsertReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ search: (param: SearchReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ query: (
|
|
|
+ param: {
|
|
|
+ collection_name: string;
|
|
|
+ } & QueryDto
|
|
|
+ ) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ deleteEntities: (param: DeleteEntitiesReq) => {
|
|
|
+ const { collection_name } = param;
|
|
|
+ if (!collection_name) {
|
|
|
+ return { status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION) };
|
|
|
+ }
|
|
|
+ return {
|
|
|
+ status: mockStatusInfo(CodeEnum.success),
|
|
|
+ data: param,
|
|
|
+ };
|
|
|
+ },
|
|
|
+ },
|
|
|
+ };
|
|
|
+});
|
|
|
+
|
|
|
+export default mockMilvusClient;
|