1/**
2 * Copyright 2021 Alibaba, Inc. and its affiliates. All Rights Reserved.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 *
16 * \author guonix
17 * \date Dec 2020
18 * \brief
19 */
20#pragma once
21
22#include <gmock/gmock.h>
23#include "query/context.h"
24
25using namespace proxima::be;
26using namespace proxima::be::query;
27
28using namespace ::testing; // for testing::*
29
30using QueryRequest = proxima::be::proto::QueryRequest;
31using QueryResponse = proxima::be::proto::QueryResponse;
32
33//! QueryContext which store request and response
34class MockQueryContext : public QueryContext {
35 public:
36 ~MockQueryContext() override = default;
37
38 //! Retrieve the PB request of query
39 MOCK_METHOD(const QueryRequest *, request, (), (const, override));
40
41 //! Retrieve the PB response of query
42 MOCK_METHOD(const QueryResponse *, response, (), (const, override));
43
44 //! Retrieve the PB response of query
45 MOCK_METHOD(QueryResponse *, mutable_response, (), (override));
46};
47
48//! CollectionQueryContext provide the collection of query
49class MockCollectionQueryContext : public CollectionQueryContext {
50 public:
51 //! Destructor
52 ~MockCollectionQueryContext() override = default;
53
54 public:
55 //! Retrieve collection name
56 MOCK_METHOD(const std::string &, collection, (), (const, override));
57};
58
59//! KNNQueryContext: Provide all the params needed for invoke segment.knn_search
60class MockKNNQueryContext : public KNNQueryContext {
61 public:
62 //! Destructor
63 ~MockKNNQueryContext() override = default;
64
65 public:
66 //! Retrieve column name
67 MOCK_METHOD(const std::string &, column, (), (const, override));
68
69 //! Retrieve features field
70 MOCK_METHOD(const std::string &, features, (), (const, override));
71
72 //! Retrieve batch_count field
73 MOCK_METHOD(uint32_t, batch_count, (), (const, override));
74
75 //! Retrieve QueryParams
76 MOCK_METHOD(const index::QueryParams &, query_params, (), (const, override));
77};
78
79//! EqualQueryContext
80class MockEqualQueryContext : public QueryKeyContext {
81 public:
82 //! Destructor
83 ~MockEqualQueryContext() override = default;
84
85 public:
86 //! Retrieve primary_key
87 MOCK_METHOD(uint64_t, primary_key, (), (const, override));
88};