test_rpc.cpp 18.9 KB
Newer Older
J
jinhai 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

Y
Yu Kun 已提交
18 19
#include <gtest/gtest.h>
#include <boost/filesystem.hpp>
Y
youny626 已提交
20
#include <thread>
Y
Yu Kun 已提交
21 22 23 24 25

#include "server/Server.h"
#include "server/grpc_impl/GrpcRequestHandler.h"
#include "server/grpc_impl/GrpcRequestScheduler.h"
#include "server/grpc_impl/GrpcRequestTask.h"
Y
youny626 已提交
26
#include "src/config.h"
Y
Yu Kun 已提交
27 28 29 30 31

#include "grpc/gen-milvus/milvus.grpc.pb.h"
#include "grpc/gen-status/status.pb.h"

#include "scheduler/ResourceFactory.h"
Y
youny626 已提交
32 33 34
#include "scheduler/SchedInst.h"
#include "server/Config.h"
#include "server/DBWrapper.h"
Y
Yu Kun 已提交
35 36
#include "utils/CommonUtil.h"

S
starlord 已提交
37
namespace {
Y
Yu Kun 已提交
38

S
starlord 已提交
39
static const char* TABLE_NAME = "test_grpc";
Y
Yu Kun 已提交
40 41
static constexpr int64_t TABLE_DIM = 256;
static constexpr int64_t INDEX_FILE_SIZE = 1024;
Y
Yu Kun 已提交
42
static constexpr int64_t VECTOR_COUNT = 1000;
Y
Yu Kun 已提交
43 44 45
static constexpr int64_t INSERT_LOOP = 10;
constexpr int64_t SECONDS_EACH_HOUR = 3600;

G
groot 已提交
46 47 48 49 50 51
void CopyRowRecord(::milvus::grpc::RowRecord* target, const std::vector<float>& src) {
    auto vector_data = target->mutable_vector_data();
    vector_data->Resize(static_cast<int>(src.size()), 0.0);
    memcpy(vector_data->mutable_data(), src.data(), src.size()* sizeof(float));
}

Y
Yu Kun 已提交
52 53 54 55
class RpcHandlerTest : public testing::Test {
 protected:
    void
    SetUp() override {
S
starlord 已提交
56
        auto res_mgr = milvus::scheduler::ResMgrInst::GetInstance();
Y
Yu Kun 已提交
57
        res_mgr->Clear();
S
starlord 已提交
58 59 60
        res_mgr->Add(milvus::scheduler::ResourceFactory::Create("disk", "DISK", 0, true, false));
        res_mgr->Add(milvus::scheduler::ResourceFactory::Create("cpu", "CPU", 0, true, true));
        res_mgr->Add(milvus::scheduler::ResourceFactory::Create("gtx1660", "GPU", 0, true, true));
Y
Yu Kun 已提交
61

S
starlord 已提交
62 63
        auto default_conn = milvus::scheduler::Connection("IO", 500.0);
        auto PCIE = milvus::scheduler::Connection("IO", 11000.0);
Y
Yu Kun 已提交
64 65 66
        res_mgr->Connect("disk", "cpu", default_conn);
        res_mgr->Connect("cpu", "gtx1660", PCIE);
        res_mgr->Start();
S
starlord 已提交
67 68
        milvus::scheduler::SchedInst::GetInstance()->Start();
        milvus::scheduler::JobMgrInst::GetInstance()->Start();
Y
Yu Kun 已提交
69

S
starlord 已提交
70
        milvus::engine::DBOptions opt;
Y
Yu Kun 已提交
71

S
starlord 已提交
72 73 74 75 76 77 78
        milvus::server::Config::GetInstance().SetDBConfigBackendUrl("sqlite://:@:/");
        milvus::server::Config::GetInstance().SetDBConfigPrimaryPath("/tmp/milvus_test");
        milvus::server::Config::GetInstance().SetDBConfigSecondaryPath("");
        milvus::server::Config::GetInstance().SetDBConfigArchiveDiskThreshold("");
        milvus::server::Config::GetInstance().SetDBConfigArchiveDaysThreshold("");
        milvus::server::Config::GetInstance().SetCacheConfigCacheInsertData("");
        milvus::server::Config::GetInstance().SetEngineConfigOmpThreadNum("");
Y
Yu Kun 已提交
79

Y
youny626 已提交
80 81 82 83 84 85 86
        //        serverConfig.SetValue(server::CONFIG_CLUSTER_MODE, "cluster");
        //        DBWrapper::GetInstance().GetInstance().StartService();
        //        DBWrapper::GetInstance().GetInstance().StopService();
        //
        //        serverConfig.SetValue(server::CONFIG_CLUSTER_MODE, "read_only");
        //        DBWrapper::GetInstance().GetInstance().StartService();
        //        DBWrapper::GetInstance().GetInstance().StopService();
Y
Yu Kun 已提交
87

S
starlord 已提交
88 89
        milvus::server::Config::GetInstance().SetResourceConfigMode("single");
        milvus::server::DBWrapper::GetInstance().StartService();
Y
Yu Kun 已提交
90

Y
youny626 已提交
91
        // initialize handler, create table
S
starlord 已提交
92
        handler = std::make_shared<milvus::server::grpc::GrpcRequestHandler>();
Y
Yu Kun 已提交
93 94
        ::grpc::ServerContext context;
        ::milvus::grpc::TableSchema request;
Y
Yu Kun 已提交
95
        ::milvus::grpc::Status status;
Y
Yu Kun 已提交
96
        request.set_table_name(TABLE_NAME);
Y
Yu Kun 已提交
97 98 99 100 101 102 103 104
        request.set_dimension(TABLE_DIM);
        request.set_index_file_size(INDEX_FILE_SIZE);
        request.set_metric_type(1);
        ::grpc::Status grpc_status = handler->CreateTable(&context, &request, &status);
    }

    void
    TearDown() override {
S
starlord 已提交
105 106 107 108
        milvus::server::DBWrapper::GetInstance().StopService();
        milvus::scheduler::JobMgrInst::GetInstance()->Stop();
        milvus::scheduler::ResMgrInst::GetInstance()->Stop();
        milvus::scheduler::SchedInst::GetInstance()->Stop();
Y
Yu Kun 已提交
109 110
        boost::filesystem::remove_all("/tmp/milvus_test");
    }
S
starlord 已提交
111

Y
Yu Kun 已提交
112
 protected:
S
starlord 已提交
113
    std::shared_ptr<milvus::server::grpc::GrpcRequestHandler> handler;
Y
Yu Kun 已提交
114 115
};

S
starlord 已提交
116
void
Y
youny626 已提交
117
BuildVectors(int64_t from, int64_t to, std::vector<std::vector<float>>& vector_record_array) {
Y
Yu Kun 已提交
118 119 120 121 122 123 124 125 126
    if (to <= from) {
        return;
    }

    vector_record_array.clear();
    for (int64_t k = from; k < to; k++) {
        std::vector<float> record;
        record.resize(TABLE_DIM);
        for (int64_t i = 0; i < TABLE_DIM; i++) {
S
starlord 已提交
127
            record[i] = (float)(k % (i + 1));
Y
Yu Kun 已提交
128 129 130 131 132 133
        }

        vector_record_array.emplace_back(record);
    }
}

S
starlord 已提交
134 135
std::string
CurrentTmDate(int64_t offset_day = 0) {
Y
Yu Kun 已提交
136 137 138 139
    time_t tt;
    time(&tt);
    tt = tt + 8 * SECONDS_EACH_HOUR;
    tt = tt + 24 * SECONDS_EACH_HOUR * offset_day;
S
starlord 已提交
140 141
    tm t;
    gmtime_r(&tt, &t);
Y
Yu Kun 已提交
142

Y
youny626 已提交
143 144
    std::string str =
        std::to_string(t.tm_year + 1900) + "-" + std::to_string(t.tm_mon + 1) + "-" + std::to_string(t.tm_mday);
Y
Yu Kun 已提交
145 146 147

    return str;
}
S
starlord 已提交
148

Y
youny626 已提交
149
}  // namespace
Y
Yu Kun 已提交
150

S
starlord 已提交
151
TEST_F(RpcHandlerTest, HAS_TABLE_TEST) {
Y
Yu Kun 已提交
152 153 154 155
    ::grpc::ServerContext context;
    ::milvus::grpc::TableName request;
    ::milvus::grpc::BoolReply reply;
    ::grpc::Status status = handler->HasTable(&context, &request, &reply);
Y
Yu Kun 已提交
156 157
    request.set_table_name(TABLE_NAME);
    status = handler->HasTable(&context, &request, &reply);
Y
Yu Kun 已提交
158 159 160 161 162
    ASSERT_TRUE(status.error_code() == ::grpc::Status::OK.error_code());
    int error_code = reply.status().error_code();
    ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
}

S
starlord 已提交
163
TEST_F(RpcHandlerTest, INDEX_TEST) {
Y
Yu Kun 已提交
164 165 166
    ::grpc::ServerContext context;
    ::milvus::grpc::IndexParam request;
    ::milvus::grpc::Status response;
Y
Yu Kun 已提交
167
    ::grpc::Status grpc_status = handler->CreateIndex(&context, &request, &response);
Y
Yu Kun 已提交
168
    request.set_table_name("test1");
Y
Yu Kun 已提交
169 170
    handler->CreateIndex(&context, &request, &response);

Y
Yu Kun 已提交
171
    request.set_table_name(TABLE_NAME);
Y
Yu Kun 已提交
172 173
    handler->CreateIndex(&context, &request, &response);

Y
Yu Kun 已提交
174
    request.mutable_index()->set_index_type(1);
Y
Yu Kun 已提交
175 176
    handler->CreateIndex(&context, &request, &response);

Y
Yu Kun 已提交
177
    request.mutable_index()->set_nlist(16384);
Y
Yu Kun 已提交
178
    grpc_status = handler->CreateIndex(&context, &request, &response);
Y
Yu Kun 已提交
179 180
    ASSERT_EQ(grpc_status.error_code(), ::grpc::Status::OK.error_code());
    int error_code = response.error_code();
Y
youny626 已提交
181
    //    ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
Y
Yu Kun 已提交
182 183 184 185

    ::milvus::grpc::TableName table_name;
    ::milvus::grpc::IndexParam index_param;
    handler->DescribeIndex(&context, &table_name, &index_param);
Y
Yu Kun 已提交
186 187 188 189
    table_name.set_table_name("test4");
    handler->DescribeIndex(&context, &table_name, &index_param);
    table_name.set_table_name(TABLE_NAME);
    handler->DescribeIndex(&context, &table_name, &index_param);
Y
Yu Kun 已提交
190
    ::milvus::grpc::Status status;
Y
Yu Kun 已提交
191 192 193 194 195
    table_name.Clear();
    handler->DropIndex(&context, &table_name, &status);
    table_name.set_table_name("test5");
    handler->DropIndex(&context, &table_name, &status);
    table_name.set_table_name(TABLE_NAME);
Y
Yu Kun 已提交
196 197 198
    handler->DropIndex(&context, &table_name, &status);
}

S
starlord 已提交
199
TEST_F(RpcHandlerTest, INSERT_TEST) {
Y
Yu Kun 已提交
200 201 202
    ::grpc::ServerContext context;
    ::milvus::grpc::InsertParam request;
    ::milvus::grpc::Status response;
Y
Yu Kun 已提交
203

Y
Yu Kun 已提交
204 205 206 207
    request.set_table_name(TABLE_NAME);
    std::vector<std::vector<float>> record_array;
    BuildVectors(0, VECTOR_COUNT, record_array);
    ::milvus::grpc::VectorIds vector_ids;
S
starlord 已提交
208 209
    for (auto& record : record_array) {
        ::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
G
groot 已提交
210
        CopyRowRecord(grpc_record, record);
Y
Yu Kun 已提交
211 212 213 214 215
    }
    handler->Insert(&context, &request, &vector_ids);
    ASSERT_EQ(vector_ids.vector_id_array_size(), VECTOR_COUNT);
}

S
starlord 已提交
216
TEST_F(RpcHandlerTest, SEARCH_TEST) {
Y
Yu Kun 已提交
217 218
    ::grpc::ServerContext context;
    ::milvus::grpc::SearchParam request;
219
    ::milvus::grpc::TopKQueryResult response;
Y
Yu Kun 已提交
220 221 222
    //test null input
    handler->Search(&context, nullptr, &response);

Y
youny626 已提交
223
    // test invalid table name
Y
Yu Kun 已提交
224 225
    handler->Search(&context, &request, &response);

Y
youny626 已提交
226
    // test table not exist
Y
Yu Kun 已提交
227 228 229
    request.set_table_name("test3");
    handler->Search(&context, &request, &response);

Y
youny626 已提交
230
    // test invalid topk
Y
Yu Kun 已提交
231
    request.set_table_name(TABLE_NAME);
Y
Yu Kun 已提交
232 233
    handler->Search(&context, &request, &response);

Y
youny626 已提交
234
    // test invalid nprobe
Y
Yu Kun 已提交
235
    request.set_topk(10);
Y
Yu Kun 已提交
236 237
    handler->Search(&context, &request, &response);

Y
youny626 已提交
238
    // test empty query record array
Y
Yu Kun 已提交
239
    request.set_nprobe(32);
Y
Yu Kun 已提交
240 241
    handler->Search(&context, &request, &response);

Y
Yu Kun 已提交
242
    std::vector<std::vector<float>> record_array;
Y
Yu Kun 已提交
243
    BuildVectors(0, VECTOR_COUNT, record_array);
Y
Yu Kun 已提交
244
    ::milvus::grpc::InsertParam insert_param;
S
starlord 已提交
245 246
    for (auto& record : record_array) {
        ::milvus::grpc::RowRecord* grpc_record = insert_param.add_row_record_array();
G
groot 已提交
247
        CopyRowRecord(grpc_record, record);
Y
Yu Kun 已提交
248
    }
Y
youny626 已提交
249
    // insert vectors
Y
Yu Kun 已提交
250 251 252 253 254
    insert_param.set_table_name(TABLE_NAME);
    ::milvus::grpc::VectorIds vector_ids;
    handler->Insert(&context, &insert_param, &vector_ids);

    BuildVectors(0, 10, record_array);
S
starlord 已提交
255 256
    for (auto& record : record_array) {
        ::milvus::grpc::RowRecord* row_record = request.add_query_record_array();
G
groot 已提交
257
        CopyRowRecord(row_record, record);
Y
Yu Kun 已提交
258
    }
Y
Yu Kun 已提交
259 260
    handler->Search(&context, &request, &response);

Y
youny626 已提交
261
    // test search with range
S
starlord 已提交
262
    ::milvus::grpc::Range* range = request.mutable_query_range_array()->Add();
Y
Yu Kun 已提交
263 264
    range->set_start_value(CurrentTmDate(-2));
    range->set_end_value(CurrentTmDate(-3));
Y
Yu Kun 已提交
265
    handler->Search(&context, &request, &response);
Y
Yu Kun 已提交
266
    request.mutable_query_range_array()->Clear();
Y
Yu Kun 已提交
267 268 269 270

    request.set_table_name("test2");
    handler->Search(&context, &request, &response);
    request.set_table_name(TABLE_NAME);
Y
Yu Kun 已提交
271 272
    handler->Search(&context, &request, &response);

Y
Yu Kun 已提交
273
    ::milvus::grpc::SearchInFilesParam search_in_files_param;
S
starlord 已提交
274
    std::string* file_id = search_in_files_param.add_file_id_array();
Y
Yu Kun 已提交
275
    *file_id = "test_tbl";
Y
Yu Kun 已提交
276 277 278
    handler->SearchInFiles(&context, &search_in_files_param, &response);
}

S
starlord 已提交
279
TEST_F(RpcHandlerTest, TABLES_TEST) {
Y
Yu Kun 已提交
280
    ::grpc::ServerContext context;
Y
Yu Kun 已提交
281
    ::milvus::grpc::TableSchema tableschema;
Y
Yu Kun 已提交
282
    ::milvus::grpc::Status response;
Y
Yu Kun 已提交
283 284
    std::string tablename = "tbl";

Y
youny626 已提交
285 286
    // create table test
    // test null input
Y
Yu Kun 已提交
287
    handler->CreateTable(&context, nullptr, &response);
Y
youny626 已提交
288
    // test invalid table name
Y
Yu Kun 已提交
289
    handler->CreateTable(&context, &tableschema, &response);
Y
youny626 已提交
290
    // test invalid table dimension
Y
Yu Kun 已提交
291
    tableschema.set_table_name(tablename);
Y
Yu Kun 已提交
292
    handler->CreateTable(&context, &tableschema, &response);
Y
youny626 已提交
293
    // test invalid index file size
Y
Yu Kun 已提交
294
    tableschema.set_dimension(TABLE_DIM);
Y
youny626 已提交
295 296
    //    handler->CreateTable(&context, &tableschema, &response);
    // test invalid index metric type
Y
Yu Kun 已提交
297 298
    tableschema.set_index_file_size(INDEX_FILE_SIZE);
    handler->CreateTable(&context, &tableschema, &response);
Y
youny626 已提交
299
    // test table already exist
Y
Yu Kun 已提交
300 301 302
    tableschema.set_metric_type(1);
    handler->CreateTable(&context, &tableschema, &response);

Y
youny626 已提交
303 304
    // describe table test
    // test invalid table name
Y
Yu Kun 已提交
305 306 307 308 309 310 311 312 313
    ::milvus::grpc::TableName table_name;
    ::milvus::grpc::TableSchema table_schema;
    handler->DescribeTable(&context, &table_name, &table_schema);

    table_name.set_table_name(TABLE_NAME);
    ::grpc::Status status = handler->DescribeTable(&context, &table_name, &table_schema);
    ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());

    ::milvus::grpc::InsertParam request;
Y
Yu Kun 已提交
314 315 316
    std::vector<std::vector<float>> record_array;
    BuildVectors(0, VECTOR_COUNT, record_array);
    ::milvus::grpc::VectorIds vector_ids;
Y
youny626 已提交
317 318
    // Insert vectors
    // test invalid table name
Y
Yu Kun 已提交
319 320
    handler->Insert(&context, &request, &vector_ids);
    request.set_table_name(tablename);
Y
youny626 已提交
321
    // test empty row record
Y
Yu Kun 已提交
322 323
    handler->Insert(&context, &request, &vector_ids);

S
starlord 已提交
324 325
    for (auto& record : record_array) {
        ::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
G
groot 已提交
326
        CopyRowRecord(grpc_record, record);
Y
Yu Kun 已提交
327
    }
Y
youny626 已提交
328
    // test vector_id size not equal to row record size
Y
Yu Kun 已提交
329 330
    vector_ids.clear_vector_id_array();
    vector_ids.add_vector_id_array(1);
Y
Yu Kun 已提交
331
    handler->Insert(&context, &request, &vector_ids);
Y
Yu Kun 已提交
332

Y
youny626 已提交
333
    // normally test
Y
Yu Kun 已提交
334
    vector_ids.clear_vector_id_array();
Y
Yu Kun 已提交
335 336
    handler->Insert(&context, &request, &vector_ids);

Y
Yu Kun 已提交
337 338 339
    request.clear_row_record_array();
    vector_ids.clear_vector_id_array();
    for (uint64_t i = 0; i < 10; ++i) {
S
starlord 已提交
340
        ::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
G
groot 已提交
341
        CopyRowRecord(grpc_record, record_array[i]);
Y
Yu Kun 已提交
342 343 344
    }
    handler->Insert(&context, &request, &vector_ids);

Y
youny626 已提交
345
    // show tables
Y
Yu Kun 已提交
346 347 348 349
    ::milvus::grpc::Command cmd;
    ::milvus::grpc::TableNameList table_name_list;
    status = handler->ShowTables(&context, &cmd, &table_name_list);
    ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());
Y
Yu Kun 已提交
350

Y
youny626 已提交
351
    // Count Table
Y
Yu Kun 已提交
352
    ::milvus::grpc::TableRowCount count;
Y
Yu Kun 已提交
353 354
    table_name.Clear();
    status = handler->CountTable(&context, &table_name, &count);
Y
Yu Kun 已提交
355
    table_name.set_table_name(tablename);
Y
Yu Kun 已提交
356
    status = handler->CountTable(&context, &table_name, &count);
Y
Yu Kun 已提交
357
    ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());
Y
youny626 已提交
358
    //    ASSERT_EQ(count.table_row_count(), vector_ids.vector_id_array_size());
Y
Yu Kun 已提交
359

Y
youny626 已提交
360
    // Preload Table
Y
Yu Kun 已提交
361 362 363
    table_name.Clear();
    status = handler->PreloadTable(&context, &table_name, &response);
    table_name.set_table_name(TABLE_NAME);
Y
Yu Kun 已提交
364 365 366
    status = handler->PreloadTable(&context, &table_name, &response);
    ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());

Y
youny626 已提交
367
    // Drop table
Y
Yu Kun 已提交
368
    table_name.set_table_name("");
Y
youny626 已提交
369
    // test invalid table name
Y
Yu Kun 已提交
370
    ::grpc::Status grpc_status = handler->DropTable(&context, &table_name, &response);
Y
Yu Kun 已提交
371 372
    table_name.set_table_name(tablename);
    grpc_status = handler->DropTable(&context, &table_name, &response);
Y
Yu Kun 已提交
373 374 375 376 377
    ASSERT_EQ(grpc_status.error_code(), ::grpc::Status::OK.error_code());
    int error_code = status.error_code();
    ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
}

G
groot 已提交
378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415
TEST_F(RpcHandlerTest, PARTITION_TEST) {
    ::grpc::ServerContext context;
    ::milvus::grpc::TableSchema table_schema;
    ::milvus::grpc::Status response;
    std::string str_table_name = "tbl_partition";
    table_schema.set_table_name(str_table_name);
    table_schema.set_dimension(TABLE_DIM);
    table_schema.set_index_file_size(INDEX_FILE_SIZE);
    table_schema.set_metric_type(1);
    handler->CreateTable(&context, &table_schema, &response);

    ::milvus::grpc::PartitionParam partition_param;
    partition_param.set_table_name(str_table_name);
    std::string partition_name = "tbl_partition_0";
    partition_param.set_partition_name(partition_name);
    std::string partition_tag = "0";
    partition_param.set_tag(partition_tag);
    handler->CreatePartition(&context, &partition_param, &response);
    ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code());

    ::milvus::grpc::TableName table_name;
    table_name.set_table_name(str_table_name);
    ::milvus::grpc::PartitionList partition_list;
    handler->ShowPartitions(&context, &table_name, &partition_list);
    ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code());
    ASSERT_EQ(partition_list.partition_array_size(), 1);

    ::milvus::grpc::PartitionParam partition_parm;
    partition_parm.set_table_name(str_table_name);
    partition_parm.set_tag(partition_tag);
    handler->DropPartition(&context, &partition_parm, &response);
    ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code());

    partition_parm.set_partition_name(partition_name);
    handler->DropPartition(&context, &partition_parm, &response);
    ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code());
}

S
starlord 已提交
416
TEST_F(RpcHandlerTest, CMD_TEST) {
Y
Yu Kun 已提交
417 418 419 420 421
    ::grpc::ServerContext context;
    ::milvus::grpc::Command command;
    command.set_cmd("version");
    ::milvus::grpc::StringReply reply;
    handler->Cmd(&context, &command, &reply);
Y
Yu Kun 已提交
422 423
    ASSERT_EQ(reply.string_reply(), MILVUS_VERSION);

Y
Yu Kun 已提交
424 425 426 427
    command.set_cmd("tasktable");
    handler->Cmd(&context, &command, &reply);
    command.set_cmd("test");
    handler->Cmd(&context, &command, &reply);
Y
Yu Kun 已提交
428 429
}

S
starlord 已提交
430
TEST_F(RpcHandlerTest, DELETE_BY_RANGE_TEST) {
Y
Yu Kun 已提交
431
    ::grpc::ServerContext context;
G
groot 已提交
432
    ::milvus::grpc::DeleteByDateParam request;
Y
Yu Kun 已提交
433
    ::milvus::grpc::Status status;
G
groot 已提交
434 435
    handler->DeleteByDate(&context, nullptr, &status);
    handler->DeleteByDate(&context, &request, &status);
Y
Yu Kun 已提交
436

Y
Yu Kun 已提交
437
    request.set_table_name(TABLE_NAME);
S
starlord 已提交
438 439
    request.mutable_range()->set_start_value(CurrentTmDate(-3));
    request.mutable_range()->set_end_value(CurrentTmDate(-2));
Y
Yu Kun 已提交
440

G
groot 已提交
441
    ::grpc::Status grpc_status = handler->DeleteByDate(&context, &request, &status);
Y
Yu Kun 已提交
442
    int error_code = status.error_code();
Y
youny626 已提交
443
    //    ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
Y
Yu Kun 已提交
444

Y
Yu Kun 已提交
445
    request.mutable_range()->set_start_value("test6");
G
groot 已提交
446
    grpc_status = handler->DeleteByDate(&context, &request, &status);
Y
Yu Kun 已提交
447
    request.mutable_range()->set_start_value(CurrentTmDate(-2));
Y
Yu Kun 已提交
448
    request.mutable_range()->set_end_value("test6");
G
groot 已提交
449
    grpc_status = handler->DeleteByDate(&context, &request, &status);
Y
Yu Kun 已提交
450
    request.mutable_range()->set_end_value(CurrentTmDate(-2));
G
groot 已提交
451
    grpc_status = handler->DeleteByDate(&context, &request, &status);
Y
Yu Kun 已提交
452 453 454
}

//////////////////////////////////////////////////////////////////////
S
starlord 已提交
455
namespace {
S
starlord 已提交
456
class DummyTask : public milvus::server::grpc::GrpcBaseTask {
S
starlord 已提交
457
 public:
S
starlord 已提交
458
    milvus::Status
Y
Yu Kun 已提交
459
    OnExecute() override {
S
starlord 已提交
460
        return milvus::Status::OK();
Y
Yu Kun 已提交
461 462
    }

S
starlord 已提交
463
    static milvus::server::grpc::BaseTaskPtr
S
starlord 已提交
464
    Create(std::string& dummy) {
S
starlord 已提交
465
        return std::shared_ptr<milvus::server::grpc::GrpcBaseTask>(new DummyTask(dummy));
Y
Yu Kun 已提交
466 467
    }

S
starlord 已提交
468
 public:
S
starlord 已提交
469
    explicit DummyTask(std::string& dummy) : GrpcBaseTask(dummy) {
Y
Yu Kun 已提交
470 471 472 473
    }
};

class RpcSchedulerTest : public testing::Test {
S
starlord 已提交
474
 protected:
Y
Yu Kun 已提交
475 476 477 478 479 480 481 482 483
    void
    SetUp() override {
        std::string dummy = "dql";
        task_ptr = std::make_shared<DummyTask>(dummy);
    }

    std::shared_ptr<DummyTask> task_ptr;
};

Y
youny626 已提交
484
}  // namespace
S
starlord 已提交
485

S
starlord 已提交
486
TEST_F(RpcSchedulerTest, BASE_TASK_TEST) {
S
starlord 已提交
487 488
    auto status = task_ptr->Execute();
    ASSERT_TRUE(status.ok());
Y
Yu Kun 已提交
489

S
starlord 已提交
490
    milvus::server::grpc::GrpcRequestScheduler::GetInstance().Start();
Y
Yu Kun 已提交
491 492
    ::milvus::grpc::Status grpc_status;
    std::string dummy = "dql";
S
starlord 已提交
493 494
    milvus::server::grpc::BaseTaskPtr base_task_ptr = DummyTask::Create(dummy);
    milvus::server::grpc::GrpcRequestScheduler::GetInstance().ExecTask(base_task_ptr, &grpc_status);
Y
Yu Kun 已提交
495

S
starlord 已提交
496
    milvus::server::grpc::GrpcRequestScheduler::GetInstance().ExecuteTask(task_ptr);
Y
Yu Kun 已提交
497
    task_ptr = nullptr;
S
starlord 已提交
498
    milvus::server::grpc::GrpcRequestScheduler::GetInstance().ExecuteTask(task_ptr);
Y
Yu Kun 已提交
499

S
starlord 已提交
500
    milvus::server::grpc::GrpcRequestScheduler::GetInstance().Stop();
Y
Yu Kun 已提交
501
}