common_meta/kv_backend/
test_util.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::sync::Arc;
17
18use derive_builder::Builder;
19
20use crate::error::Result;
21use crate::kv_backend::txn::{Txn, TxnResponse};
22use crate::kv_backend::{
23    BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, BatchPutRequest,
24    BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, KvBackend, PutRequest, PutResponse,
25    RangeRequest, RangeResponse, TxnService,
26};
27
28pub type MockFn<Req, Resp> = Arc<dyn Fn(Req) -> Result<Resp> + Send + Sync>;
29
30/// A mock kv backend for testing.
31#[derive(Builder)]
32pub struct MockKvBackend {
33    #[builder(setter(strip_option), default)]
34    pub range_fn: Option<MockFn<RangeRequest, RangeResponse>>,
35    #[builder(setter(strip_option), default)]
36    pub put_fn: Option<MockFn<PutRequest, PutResponse>>,
37    #[builder(setter(strip_option), default)]
38    pub batch_put_fn: Option<MockFn<BatchPutRequest, BatchPutResponse>>,
39    #[builder(setter(strip_option), default)]
40    pub batch_get_fn: Option<MockFn<BatchGetRequest, BatchGetResponse>>,
41    #[builder(setter(strip_option), default)]
42    pub delete_range_fn: Option<MockFn<DeleteRangeRequest, DeleteRangeResponse>>,
43    #[builder(setter(strip_option), default)]
44    pub batch_delete_fn: Option<MockFn<BatchDeleteRequest, BatchDeleteResponse>>,
45    #[builder(setter(strip_option), default)]
46    pub txn: Option<MockFn<Txn, TxnResponse>>,
47    #[builder(setter(strip_option), default)]
48    pub max_txn_ops: Option<usize>,
49}
50
51#[async_trait::async_trait]
52impl TxnService for MockKvBackend {
53    type Error = crate::error::Error;
54
55    async fn txn(&self, txn: Txn) -> Result<TxnResponse> {
56        if let Some(f) = &self.txn {
57            f(txn)
58        } else {
59            unimplemented!()
60        }
61    }
62
63    fn max_txn_ops(&self) -> usize {
64        self.max_txn_ops.unwrap()
65    }
66}
67
68#[async_trait::async_trait]
69impl KvBackend for MockKvBackend {
70    fn name(&self) -> &str {
71        "mock_kv_backend"
72    }
73
74    fn as_any(&self) -> &dyn Any {
75        self
76    }
77
78    async fn range(&self, req: RangeRequest) -> Result<RangeResponse> {
79        if let Some(f) = &self.range_fn {
80            f(req)
81        } else {
82            unimplemented!()
83        }
84    }
85
86    async fn put(&self, req: PutRequest) -> Result<PutResponse> {
87        if let Some(f) = &self.put_fn {
88            f(req)
89        } else {
90            unimplemented!()
91        }
92    }
93
94    async fn batch_put(&self, req: BatchPutRequest) -> Result<BatchPutResponse> {
95        if let Some(f) = &self.batch_put_fn {
96            f(req)
97        } else {
98            unimplemented!()
99        }
100    }
101
102    async fn batch_get(&self, req: BatchGetRequest) -> Result<BatchGetResponse> {
103        if let Some(f) = &self.batch_get_fn {
104            f(req)
105        } else {
106            unimplemented!()
107        }
108    }
109
110    async fn delete_range(&self, req: DeleteRangeRequest) -> Result<DeleteRangeResponse> {
111        if let Some(f) = &self.delete_range_fn {
112            f(req)
113        } else {
114            unimplemented!()
115        }
116    }
117
118    async fn batch_delete(&self, req: BatchDeleteRequest) -> Result<BatchDeleteResponse> {
119        if let Some(f) = &self.batch_delete_fn {
120            f(req)
121        } else {
122            unimplemented!()
123        }
124    }
125}
common_meta/kv_backend/test_util.rs

common_meta/kv_backend/
test_util.rs
