Browse Source

Add collection service tests

tumao 3 years ago
parent
commit
20934e4461

+ 155 - 34
express/src/__mocks__/milvus/milvusClient.ts

@@ -1,12 +1,32 @@
 import {
+  AlterAliasReq,
+  CreateAliasReq,
   CreateCollectionReq,
   DescribeCollectionReq,
+  DropAliasReq,
   DropCollectionReq,
   FlushReq,
+  GetCollectionStatisticsReq,
+  GetIndexStateReq,
+  InsertReq,
   LoadCollectionReq,
-  ShowCollectionsReq
+  ReleaseLoadCollectionReq,
+  SearchReq,
+  ShowCollectionsReq,
 } from '@zilliz/milvus2-sdk-node/dist/milvus/types';
-import { CodeEnum, ERR_NO_ADDRESS, ERR_NO_COLLECTION, ERR_NO_PARAM } from '../../__tests__/utils/constants';
+import { QueryDto } from '../../collections/dto';
+import {
+  CodeEnum,
+  ERR_NO_ADDRESS,
+  ERR_NO_ALIAS,
+  ERR_NO_COLLECTION,
+  ERR_NO_PARAM,
+  mockCollectionNames,
+  mockCollections,
+  mockIndexState,
+  mockLoadedCollections,
+} from '../../__tests__/utils/constants';
+import { mockStatusInfo } from '../../__tests__/utils/mock.util';
 
 const mockMilvusClient = jest.fn().mockImplementation((address: string) => {
   return {
@@ -21,65 +41,129 @@ const mockMilvusClient = jest.fn().mockImplementation((address: string) => {
       showCollections: (param?: ShowCollectionsReq) => {
         if (!param) {
           return {
-            status: {
-              error_code: CodeEnum.error,
-              reason: ERR_NO_PARAM,
-            },
+            status: mockStatusInfo(CodeEnum.success),
+            data: mockCollectionNames,
           };
         }
-        const { collection_names } = param;
-        return {
-          status: { error_code: CodeEnum.success },
-          data: collection_names,
-        };
+        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 {
-            error_code: CodeEnum.error,
-            reason: ERR_NO_COLLECTION,
-          };
+          return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
         }
-        return { error_code: CodeEnum.success, data: fields };
+        return { ...mockStatusInfo(CodeEnum.success), data: fields };
       },
       describeCollection: (param: DescribeCollectionReq) => {
         const { collection_name } = param;
         if (!collection_name) {
           return {
-            status: {
-              error_code: CodeEnum.error,
-              reason: ERR_NO_COLLECTION,
-            },
+            status: mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION),
           };
         }
+        const res =
+          mockCollections.find((c) => c.name === collection_name) || {};
         return {
-          status: { error_code: CodeEnum.success },
-          data: collection_name,
+          status: mockStatusInfo(CodeEnum.success),
+          ...res,
         };
       },
       dropCollection: (param: DropCollectionReq) => {
         const { collection_name } = param;
         if (!collection_name) {
-          return {
-            error_code: CodeEnum.error,
-            reason: ERR_NO_COLLECTION,
-          };
+          return mockStatusInfo(CodeEnum.error, ERR_NO_COLLECTION);
         }
-        return { error_code: CodeEnum.success, data: collection_name };
+        return { ...mockStatusInfo(CodeEnum.success), data: collection_name };
       },
       loadCollection: (param: LoadCollectionReq) => {
-        const {collection_name} = param
+        const { collection_name } = param;
         if (!collection_name) {
-          return {
-            error_code: CodeEnum.error,
-            reason: ERR_NO_COLLECTION
-          }
+          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: {},
-    indexManager: {},
+    indexManager: {
+      getIndexState: (param: GetIndexStateReq) => {
+        const { collection_name } = param;
+        const data =
+          mockIndexState.find((i) => i.collection_name === collection_name) ||
+          {};
+        return {
+          ...mockStatusInfo(CodeEnum.success),
+          ...data,
+        };
+      },
+    },
     dataManager: {
       flush: (data: FlushReq) => ({
         data,
@@ -93,6 +177,43 @@ const mockMilvusClient = jest.fn().mockImplementation((address: string) => {
           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,
+        };
+      },
     },
   };
 });

+ 201 - 8
express/src/__tests__/collections/collections.service.test.ts

@@ -1,10 +1,18 @@
-
-
 import mockMilvusClient from '../../__mocks__/milvus/milvusClient';
 import { CollectionsService } from '../../collections/collections.service';
 import { MilvusService } from '../../milvus/milvus.service';
-import { ERR_NO_COLLECTION, ERR_NO_PARAM, mockAddress } from '../utils/constants';
-
+import {
+  ERR_NO_ALIAS,
+  ERR_NO_COLLECTION,
+  ERR_NO_PARAM,
+  mockAddress,
+  mockCollectionNames,
+  mockCollections,
+  mockGetAllCollectionsData,
+  mockIndexState,
+  mockLoadedCollections,
+  mockLoadedCollectionsData,
+} from '../utils/constants';
 
 // mock Milvus client
 jest.mock('@zilliz/milvus2-sdk-node', () => {
@@ -41,8 +49,14 @@ describe('Test collections service', () => {
     });
     expect(res.data.length).toBe(2);
 
+    const defaultRes = await service.getCollections();
+    expect(defaultRes.data).toEqual(mockCollectionNames);
+
+    const loadedRes = await service.getCollections({ type: 1 });
+    expect(loadedRes.data).toEqual(mockLoadedCollections);
+
     try {
-      await service.getCollections();
+      await service.getCollections({ collection_names: [] });
     } catch (err) {
       expect(err).toBe(ERR_NO_PARAM);
     }
@@ -66,7 +80,9 @@ describe('Test collections service', () => {
     const res = await service.describeCollection({
       collection_name: 'c1',
     });
-    expect(res.data).toBe('c1');
+    const { status, ...result } = res;
+    const [mockRes] = mockCollections;
+    expect(result).toEqual(mockRes);
 
     try {
       await service.describeCollection({ collection_name: '' });
@@ -87,6 +103,183 @@ describe('Test collections service', () => {
   });
 
   test('test loadCollection method', async () => {
-    
-  })
+    const res = await service.loadCollection({ collection_name: 'c1' });
+    expect(res.data).toBe('c1');
+
+    try {
+      await service.loadCollection({ collection_name: '' });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_COLLECTION);
+    }
+  });
+
+  test('test releaseCollection method', async () => {
+    const res = await service.releaseCollection({ collection_name: 'c1' });
+    expect(res.data).toBe('c1');
+
+    try {
+      await service.releaseCollection({ collection_name: '' });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_COLLECTION);
+    }
+  });
+
+  test('test getCollectionStatistics method', async () => {
+    const res = await service.getCollectionStatistics({
+      collection_name: 'c1',
+    });
+    const { status, ...data } = res;
+    expect(data.name).toBe('c1');
+    expect(data.stats.length).toBe(1);
+
+    try {
+      await service.getCollectionStatistics({ collection_name: '' });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_COLLECTION);
+    }
+  });
+
+  test('test insert method', async () => {
+    const mockParam = {
+      collection_name: 'c1',
+      fields_data: {
+        vector_field: [1, 2, 3, 4],
+        age: 7,
+      },
+    };
+    const res = await service.insert(mockParam);
+    expect(res.data).toEqual(mockParam);
+
+    try {
+      await service.insert({
+        collection_name: '',
+        fields_data: {
+          vector_field: [1, 2, 3, 4],
+        },
+      });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_COLLECTION);
+    }
+  });
+
+  test('test vectorSearch method', async () => {
+    const mockParam = {
+      collection_name: 'c1',
+      search_params: {
+        anns_field: 'float_vector',
+        topk: '10',
+        metric_type: 'L2',
+        params: JSON.stringify({ nprobe: 1024 }),
+      },
+      vectors: [[1, 2, 3, 4]],
+      vector_type: 101,
+    };
+
+    const res = await service.vectorSearch(mockParam);
+    expect(res.data).toEqual(mockParam);
+
+    try {
+      await service.vectorSearch({ ...mockParam, collection_name: '' });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_COLLECTION);
+    }
+  });
+
+  test('test createAlias method', async () => {
+    const mockParam = {
+      collection_name: 'c1',
+      alias: 'alias',
+    };
+
+    const res = await service.createAlias(mockParam);
+    expect(res.data).toEqual(mockParam);
+
+    try {
+      await service.createAlias({ collection_name: '', alias: '' });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_COLLECTION);
+    }
+  });
+
+  test('test alterAlias method', async () => {
+    const mockParam = {
+      collection_name: 'c1',
+      alias: 'alias',
+    };
+
+    const res = await service.alterAlias(mockParam);
+    expect(res.data).toEqual(mockParam);
+
+    try {
+      await service.alterAlias({ collection_name: '', alias: '' });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_COLLECTION);
+    }
+  });
+
+  test('test dropAlias method', async () => {
+    const res = await service.dropAlias({ alias: 'alias' });
+    expect(res.data).toBe('alias');
+
+    try {
+      await service.dropAlias({ alias: '' });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_ALIAS);
+    }
+  });
+
+  test('test query method', async () => {
+    const mockParam = {
+      collection_name: 'c1',
+      expr: 'age > 7',
+    };
+    const res = await service.query(mockParam);
+    expect(res.data).toEqual(mockParam);
+
+    try {
+      await service.query({ collection_name: '', expr: '' });
+    } catch (err) {
+      expect(err).toBe(ERR_NO_COLLECTION);
+    }
+  });
+
+  test('test getIndexStatus method', async () => {
+    const res = await service.getIndexStatus({ collection_name: 'c1' });
+    const { error_code, ...data } = res;
+    expect(data).toEqual({ collection_name: 'c1', state: 3 });
+  });
+
+  test('test getAllCollections method', async () => {
+    const res = await service.getAllCollections();
+    expect(res).toEqual(mockGetAllCollectionsData);
+  });
+
+  test('test getLoadedCollections method', async () => {
+    const res = await service.getLoadedColletions();
+    expect(res).toEqual(mockLoadedCollectionsData);
+  });
+
+  test('test getStatistics method', async () => {
+    const res = await service.getStatistics();
+    expect(res).toEqual({
+      // 2 collections
+      collectionCount: 2,
+      // each collection 7 row counts
+      totalData: 14,
+    });
+  });
+
+  test('test getCollectionIndexStatus method', async () => {
+    const res = await service.getCollectionsIndexStatus();
+    expect(res).toEqual([
+      {
+        collection_name: 'c1',
+        index_status: 3,
+      },
+      {
+        collection_name: 'c2',
+        index_status: 2,
+      },
+    ]);
+  });
 });

+ 4 - 3
express/src/__tests__/milvus/index.test.ts → express/src/__tests__/milvus/milvus.controller.test.ts

@@ -1,11 +1,12 @@
 import express from 'express';
 import http from 'http';
 import supertest from 'supertest';
+import { TransformResMiddlerware } from '../../middlewares';
 import { router as connectRouter } from '../../milvus/index';
-import { TransformResInterceptor } from '../../interceptors';
+
 import MilvusService from '../../__mocks__/milvus/milvus.service';
+import { mockAddress } from '../utils/constants';
 
-const mockAddress = '127.0.0.1';
 
 // mock Milvus client service
 jest.mock('../../__mocks__/milvus/milvus.service');
@@ -20,7 +21,7 @@ describe('Test Milvus Module', () => {
     app = express();
     const router = express.Router();
     router.use('/milvus', connectRouter);
-    app.use(TransformResInterceptor);
+    app.use(TransformResMiddlerware);
     app.use('/api/v1', router);
     server = http.createServer(app);
     server.listen(done);

+ 152 - 5
express/src/__tests__/utils/constants.ts

@@ -1,12 +1,159 @@
 export enum CodeEnum {
   success = 'Success',
-  error = 'Error'
+  error = 'Error',
 }
 
+export type CodeStatus = {
+  error_code: CodeEnum;
+  reason?: string;
+};
+
 // error msgs
-export const ERR_NO_COLLECTION = 'collection name is invalid'
-export const ERR_NO_ADDRESS = 'no address'
-export const ERR_NO_PARAM = 'no valid param'
+export const ERR_NO_COLLECTION = 'collection name is invalid';
+export const ERR_NO_ADDRESS = 'no address';
+export const ERR_NO_PARAM = 'no valid param';
+export const ERR_NO_ALIAS = 'no valid alias';
 
 // mock data
-export const mockAddress = '127.0.0.1'
+export const mockAddress = '127.0.0.1';
+export const mockCollectionNames = [{ name: 'c1' }, { name: 'c2' }];
+export const mockCollections = [
+  {
+    name: 'c1',
+    collectionID: 1,
+    schema: {
+      fields: [
+        {
+          name: 'vector_field',
+          data_type: 'data_type',
+          type_params: [
+            {
+              key: 'dim',
+              value: '4',
+            },
+          ],
+        },
+        {
+          is_primary_key: true,
+          autoID: true,
+          name: 'age',
+          data_type: 'data_type',
+          type_params: [] as any[],
+        },
+      ],
+      description: 'mock schema description 1',
+    },
+    created_utc_timestamp: '123456',
+  },
+  {
+    name: 'c2',
+    collectionID: 2,
+    schema: {
+      fields: [
+        {
+          name: 'vector_field',
+          data_type: 'data_type',
+          type_params: [
+            {
+              key: 'dim',
+              value: '4',
+            },
+          ],
+        },
+        {
+          name: 'age',
+          data_type: 'data_type',
+          type_params: [] as any[],
+        },
+      ],
+      description: 'mock schema description 2',
+    },
+    created_utc_timestamp: '1234567',
+  },
+];
+export const mockLoadedCollections = [
+  {
+    id: 1,
+    name: 'c1',
+    loadedPercentage: '100',
+  },
+];
+// index state is finished
+export const mockIndexState = [
+  { collection_name: 'c1', state: 3 },
+  { collection_name: 'c2', state: 2 },
+];
+
+// mock results
+export const mockGetAllCollectionsData = [
+  {
+    collection_name: 'c2',
+    schema: {
+      fields: [
+        {
+          name: 'vector_field',
+          data_type: 'data_type',
+          type_params: [
+            {
+              key: 'dim',
+              value: '4',
+            },
+          ],
+        },
+        {
+          name: 'age',
+          data_type: 'data_type',
+          type_params: [] as any[],
+        },
+      ],
+      description: 'mock schema description 2',
+    },
+    description: 'mock schema description 2',
+    autoID: undefined as boolean,
+    rowCount: 7,
+    id: 2,
+    loadedPercentage: '-1',
+    createdTime: 1234567,
+    index_status: 2,
+  },
+  {
+    collection_name: 'c1',
+    schema: {
+      fields: [
+        {
+          name: 'vector_field',
+          data_type: 'data_type',
+          type_params: [
+            {
+              key: 'dim',
+              value: '4',
+            },
+          ],
+        },
+        {
+          is_primary_key: true,
+          autoID: true,
+          name: 'age',
+          data_type: 'data_type',
+          type_params: [] as any[],
+        },
+      ],
+      description: 'mock schema description 1',
+    },
+    description: 'mock schema description 1',
+    autoID: true,
+    rowCount: 7,
+    id: 1,
+    loadedPercentage: '100',
+    createdTime: 123456,
+    index_status: 3,
+  },
+];
+
+export const mockLoadedCollectionsData = [
+  {
+    collection_name: 'c1',
+    id: 1,
+    rowCount: 7,
+  },
+];

+ 13 - 0
express/src/__tests__/utils/mock.util.ts

@@ -0,0 +1,13 @@
+import { CodeEnum, CodeStatus } from './constants';
+
+export const mockStatusInfo = (type: CodeEnum, msg?: string): CodeStatus => {
+  if (type === CodeEnum.success) {
+    return {
+      error_code: type,
+    };
+  }
+  return {
+    error_code: type,
+    reason: msg,
+  };
+};

+ 2 - 3
express/src/collections/collections.service.ts

@@ -19,8 +19,8 @@ import {
   DropAliasReq,
   ShowCollectionsReq,
   ShowCollectionsType,
-} from "@zilliz/milvus2-sdk-node/dist/milvus/types/Collection";
-import { QueryDto } from "./dto";
+} from '@zilliz/milvus2-sdk-node/dist/milvus/types/Collection';
+import { QueryDto } from './dto';
 
 export class CollectionsService {
   constructor(private milvusService: MilvusService) {}
@@ -144,7 +144,6 @@ export class CollectionsService {
     if (res.data.length > 0) {
       for (const item of res.data) {
         const { name } = item;
-
         const collectionInfo = await this.describeCollection({
           collection_name: name,
         });