mirror of https://github.com/milvus-io/milvus.git
Update test and some interface
parent
d7e75bf95b
commit
3da09021e6
|
@ -0,0 +1,17 @@
|
|||
import pytest
|
||||
from flask import Flask
|
||||
from engine import app
|
||||
|
||||
@pytest.fixture(scope='module')
|
||||
def test_client():
|
||||
# Flask provides a way to test your application by exposing the Werkzeug test Client
|
||||
# and handling the context locals for you.
|
||||
testing_client = app.test_client()
|
||||
|
||||
# Establish an application context before running the tests.
|
||||
ctx = app.app_context()
|
||||
ctx.push()
|
||||
|
||||
yield testing_client # this is where the testing happens!
|
||||
|
||||
ctx.pop()
|
|
@ -0,0 +1,36 @@
|
|||
from engine.controller.group_handler import GroupHandler
|
||||
from engine.settings import DATABASE_DIRECTORY
|
||||
import pytest
|
||||
import os
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class TestGroupHandler:
|
||||
def test_get_group(self):
|
||||
group_path = GroupHandler.GetGroupDirectory('test_group')
|
||||
verified_path = DATABASE_DIRECTORY + '/' + 'test_group'
|
||||
logger.debug(group_path)
|
||||
assert group_path == verified_path
|
||||
|
||||
def test_create_group(self):
|
||||
group_path = GroupHandler.CreateGroupDirectory('test_group')
|
||||
if os.path.exists(group_path):
|
||||
assert True
|
||||
else:
|
||||
assert False
|
||||
|
||||
def test_delete_group(self):
|
||||
group_path = GroupHandler.GetGroupDirectory('test_group')
|
||||
if os.path.exists(group_path):
|
||||
assert True
|
||||
GroupHandler.DeleteGroupDirectory('test_group')
|
||||
if os.path.exists(group_path):
|
||||
assert False
|
||||
else:
|
||||
assert True
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
|
@ -0,0 +1,102 @@
|
|||
from engine.controller.vector_engine import VectorEngine
|
||||
from engine.settings import DATABASE_DIRECTORY
|
||||
from flask import jsonify
|
||||
import pytest
|
||||
import os
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class TestVectorEngine:
|
||||
def setup_class(self):
|
||||
self.__vector = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]
|
||||
self.__limit = 3
|
||||
|
||||
|
||||
def teardown_class(self):
|
||||
pass
|
||||
|
||||
def test_group(self):
|
||||
# Make sure there is no group
|
||||
code, group_id, file_number = VectorEngine.DeleteGroup('test_group')
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
assert group_id == 'test_group'
|
||||
assert file_number == 0
|
||||
|
||||
# Add a group
|
||||
code, group_id, file_number = VectorEngine.AddGroup('test_group', 8)
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
assert group_id == 'test_group'
|
||||
assert file_number == 0
|
||||
|
||||
# Check the group existing
|
||||
code, group_id, file_number = VectorEngine.GetGroup('test_group')
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
assert group_id == 'test_group'
|
||||
assert file_number == 0
|
||||
|
||||
# Check the group list
|
||||
code, group_list = VectorEngine.GetGroupList()
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
assert group_list == [{'group_name': 'test_group', 'file_number': 0}]
|
||||
|
||||
# Add Vector for not exist group
|
||||
code = VectorEngine.AddVector('not_exist_group', self.__vector)
|
||||
assert code == VectorEngine.GROUP_NOT_EXIST
|
||||
|
||||
# Add vector for exist group
|
||||
code = VectorEngine.AddVector('test_group', self.__vector)
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
|
||||
# Check search vector interface
|
||||
code, vector_id = VectorEngine.SearchVector('test_group', self.__vector, self.__limit)
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
assert vector_id == 0
|
||||
|
||||
# Check create index interface
|
||||
code = VectorEngine.CreateIndex('test_group')
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
|
||||
# Remove the group
|
||||
code, group_id, file_number = VectorEngine.DeleteGroup('test_group')
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
assert group_id == 'test_group'
|
||||
assert file_number == 0
|
||||
|
||||
# Check the group is disppeared
|
||||
code, group_id, file_number = VectorEngine.GetGroup('test_group')
|
||||
assert code == VectorEngine.FAULT_CODE
|
||||
assert group_id == 'test_group'
|
||||
assert file_number == 0
|
||||
|
||||
# Check SearchVector interface
|
||||
code = VectorEngine.SearchVector('test_group', self.__vector, self.__limit)
|
||||
assert code == VectorEngine.GROUP_NOT_EXIST
|
||||
|
||||
# Create Index for not exist group id
|
||||
code = VectorEngine.CreateIndex('test_group')
|
||||
assert code == VectorEngine.GROUP_NOT_EXIST
|
||||
|
||||
# Clear raw file
|
||||
code = VectorEngine.ClearRawFile('test_group')
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
|
||||
def test_raw_file(self):
|
||||
filename = VectorEngine.InsertVectorIntoRawFile('test_group', 'test_group.raw', self.__vector)
|
||||
assert filename == 'test_group.raw'
|
||||
|
||||
expected_list = [self.__vector]
|
||||
vector_list = VectorEngine.GetVectorListFromRawFile('test_group', filename)
|
||||
|
||||
print('expected_list: ', expected_list)
|
||||
print('vector_list: ', vector_list)
|
||||
|
||||
assert vector_list == expected_list
|
||||
|
||||
code = VectorEngine.ClearRawFile('test_group')
|
||||
assert code == VectorEngine.SUCCESS_CODE
|
||||
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,71 @@
|
|||
from engine.controller.vector_engine import VectorEngine
|
||||
from engine.settings import DATABASE_DIRECTORY
|
||||
from engine import app
|
||||
from flask import jsonify
|
||||
import pytest
|
||||
import os
|
||||
import logging
|
||||
import json
|
||||
|
||||
logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class TestViews:
|
||||
|
||||
def loads(self, resp):
|
||||
return json.loads(resp.data.decode())
|
||||
|
||||
def test_group(self, test_client):
|
||||
data = {"dimension": 10}
|
||||
|
||||
resp = test_client.get('/vector/group/6')
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 1
|
||||
|
||||
resp = test_client.post('/vector/group/6', data=json.dumps(data))
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
|
||||
resp = test_client.get('/vector/group/6')
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
|
||||
# GroupList
|
||||
resp = test_client.get('/vector/group')
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
assert self.loads(resp)['group_list'] == [{'file_number': 0, 'group_name': '6'}]
|
||||
|
||||
resp = test_client.delete('/vector/group/6')
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
|
||||
|
||||
def test_vector(self, test_client):
|
||||
dimension = {"dimension": 10}
|
||||
resp = test_client.post('/vector/group/6', data=json.dumps(dimension))
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
|
||||
vector = {"vector": [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8]}
|
||||
resp = test_client.post('/vector/add/6', data=json.dumps(vector))
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
|
||||
resp = test_client.post('/vector/index/6')
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
|
||||
limit = {"vector": [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8], "limit": 3}
|
||||
resp = test_client.get('/vector/search/6', data=json.dumps(limit))
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
assert self.loads(resp)['vector_id'] == 0
|
||||
|
||||
resp = test_client.delete('/vector/group/6')
|
||||
assert resp.status_code == 200
|
||||
assert self.loads(resp)['code'] == 0
|
||||
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,84 @@
|
|||
from flask import Flask, jsonify, request
|
||||
from flask_restful import Resource, Api
|
||||
from engine import app, db
|
||||
from engine.model.group_table import GroupTable
|
||||
from engine.controller.vector_engine import VectorEngine
|
||||
|
||||
# app = Flask(__name__)
|
||||
api = Api(app)
|
||||
|
||||
|
||||
from flask_restful import reqparse
|
||||
from flask_restful import request
|
||||
class Vector(Resource):
|
||||
def __init__(self):
|
||||
self.__parser = reqparse.RequestParser()
|
||||
self.__parser.add_argument('vector', type=float, action='append', location=['json'])
|
||||
|
||||
def post(self, group_id):
|
||||
args = self.__parser.parse_args()
|
||||
vector = args['vector']
|
||||
code = VectorEngine.AddVector(group_id, vector)
|
||||
return jsonify({'code': code})
|
||||
|
||||
|
||||
class VectorSearch(Resource):
|
||||
def __init__(self):
|
||||
self.__parser = reqparse.RequestParser()
|
||||
self.__parser.add_argument('vector', type=float, action='append', location=['json'])
|
||||
self.__parser.add_argument('limit', type=int, action='append', location=['json'])
|
||||
|
||||
def get(self, group_id):
|
||||
args = self.__parser.parse_args()
|
||||
print('vector: ', args['vector'])
|
||||
# go to search every thing
|
||||
code, vector_id = VectorEngine.SearchVector(group_id, args['vector'], args['limit'])
|
||||
return jsonify({'code': code, 'vector_id': vector_id})
|
||||
|
||||
|
||||
class Index(Resource):
|
||||
def __init__(self):
|
||||
self.__parser = reqparse.RequestParser()
|
||||
# self.__parser.add_argument('group_id', type=str)
|
||||
|
||||
def post(self, group_id):
|
||||
code = VectorEngine.CreateIndex(group_id)
|
||||
return jsonify({'code': code})
|
||||
|
||||
|
||||
class Group(Resource):
|
||||
def __init__(self):
|
||||
self.__parser = reqparse.RequestParser()
|
||||
self.__parser.add_argument('group_id', type=str)
|
||||
self.__parser.add_argument('dimension', type=int, action='append', location=['json'])
|
||||
|
||||
def post(self, group_id):
|
||||
args = self.__parser.parse_args()
|
||||
dimension = args['dimension']
|
||||
code, group_id, file_number = VectorEngine.AddGroup(group_id, dimension)
|
||||
return jsonify({'code': code, 'group': group_id, 'filenumber': file_number})
|
||||
|
||||
def get(self, group_id):
|
||||
code, group_id, file_number = VectorEngine.GetGroup(group_id)
|
||||
return jsonify({'code': code, 'group': group_id, 'filenumber': file_number})
|
||||
|
||||
def delete(self, group_id):
|
||||
code, group_id, file_number = VectorEngine.DeleteGroup(group_id)
|
||||
return jsonify({'code': code, 'group': group_id, 'filenumber': file_number})
|
||||
|
||||
|
||||
class GroupList(Resource):
|
||||
def get(self):
|
||||
code, group_list = VectorEngine.GetGroupList()
|
||||
return jsonify({'code': code, 'group_list': group_list})
|
||||
|
||||
|
||||
api.add_resource(Vector, '/vector/add/<group_id>')
|
||||
api.add_resource(Group, '/vector/group/<group_id>')
|
||||
api.add_resource(GroupList, '/vector/group')
|
||||
api.add_resource(Index, '/vector/index/<group_id>')
|
||||
api.add_resource(VectorSearch, '/vector/search/<group_id>')
|
||||
|
||||
|
||||
# if __name__ == '__main__':
|
||||
# app.run()
|
|
@ -0,0 +1 @@
|
|||
pytest -v --disable-warnings
|
|
@ -0,0 +1,12 @@
|
|||
import os
|
||||
import faiss
|
||||
|
||||
class StorageManager(object):
|
||||
def __init__():
|
||||
pass
|
||||
|
||||
def put(vector, directory, index_type):
|
||||
pass
|
||||
|
||||
def take(dir):
|
||||
pass
|
Loading…
Reference in New Issue