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 | |
25 | using namespace proxima::be; |
26 | using namespace proxima::be::query; |
27 | |
28 | using namespace ::testing; // for testing::* |
29 | |
30 | using QueryRequest = proxima::be::proto::QueryRequest; |
31 | using QueryResponse = proxima::be::proto::QueryResponse; |
32 | |
33 | //! QueryContext which store request and response |
34 | class 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 |
49 | class 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 |
60 | class 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 |
80 | class 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 | }; |