object_store/layers/
mock.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::fmt::Debug;
16use std::sync::Arc;
17
18use derive_builder::Builder;
19pub use oio::*;
20pub use opendal::raw::{
21    Access, Layer, LayeredAccess, OpDelete, OpList, OpRead, OpWrite, RpDelete, RpList, RpRead,
22    RpWrite, oio,
23};
24pub use opendal::{Buffer, Error, ErrorKind, Metadata, Result};
25
26pub type MockWriterFactory = Arc<dyn Fn(&str, OpWrite, oio::Writer) -> oio::Writer + Send + Sync>;
27pub type MockReaderFactory = Arc<dyn Fn(&str, OpRead, oio::Reader) -> oio::Reader + Send + Sync>;
28pub type MockListerFactory = Arc<dyn Fn(&str, OpList, oio::Lister) -> oio::Lister + Send + Sync>;
29pub type MockDeleterFactory = Arc<dyn Fn(oio::Deleter) -> oio::Deleter + Send + Sync>;
30
31#[derive(Builder)]
32pub struct MockLayer {
33    #[builder(setter(strip_option), default)]
34    writer_factory: Option<MockWriterFactory>,
35    #[builder(setter(strip_option), default)]
36    reader_factory: Option<MockReaderFactory>,
37    #[builder(setter(strip_option), default)]
38    lister_factory: Option<MockListerFactory>,
39    #[builder(setter(strip_option), default)]
40    deleter_factory: Option<MockDeleterFactory>,
41}
42
43impl Clone for MockLayer {
44    fn clone(&self) -> Self {
45        Self {
46            writer_factory: self.writer_factory.clone(),
47            reader_factory: self.reader_factory.clone(),
48            lister_factory: self.lister_factory.clone(),
49            deleter_factory: self.deleter_factory.clone(),
50        }
51    }
52}
53
54impl<A: Access> Layer<A> for MockLayer {
55    type LayeredAccess = MockAccessor<A>;
56
57    fn layer(&self, inner: A) -> Self::LayeredAccess {
58        MockAccessor {
59            inner,
60            writer_factory: self.writer_factory.clone(),
61            reader_factory: self.reader_factory.clone(),
62            lister_factory: self.lister_factory.clone(),
63            deleter_factory: self.deleter_factory.clone(),
64        }
65    }
66}
67
68pub struct MockAccessor<A> {
69    inner: A,
70    writer_factory: Option<MockWriterFactory>,
71    reader_factory: Option<MockReaderFactory>,
72    lister_factory: Option<MockListerFactory>,
73    deleter_factory: Option<MockDeleterFactory>,
74}
75
76impl<A: Debug> Debug for MockAccessor<A> {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        f.debug_struct("MockAccessor")
79            .field("inner", &self.inner)
80            .finish()
81    }
82}
83
84pub struct MockReader {
85    inner: oio::Reader,
86}
87
88impl oio::Read for MockReader {
89    async fn read(&mut self) -> Result<Buffer> {
90        self.inner.read().await
91    }
92}
93
94pub struct MockWriter {
95    inner: oio::Writer,
96}
97
98impl oio::Write for MockWriter {
99    async fn write(&mut self, bs: Buffer) -> Result<()> {
100        self.inner.write(bs).await
101    }
102
103    async fn close(&mut self) -> Result<Metadata> {
104        self.inner.close().await
105    }
106
107    async fn abort(&mut self) -> Result<()> {
108        self.inner.abort().await
109    }
110}
111
112pub struct MockLister {
113    inner: oio::Lister,
114}
115
116impl oio::List for MockLister {
117    async fn next(&mut self) -> Result<Option<oio::Entry>> {
118        self.inner.next().await
119    }
120}
121
122pub struct MockDeleter {
123    inner: oio::Deleter,
124}
125
126impl oio::Delete for MockDeleter {
127    fn delete(&mut self, path: &str, args: OpDelete) -> Result<()> {
128        self.inner.delete(path, args)
129    }
130
131    async fn flush(&mut self) -> Result<usize> {
132        self.inner.flush().await
133    }
134}
135
136impl<A: Access> LayeredAccess for MockAccessor<A> {
137    type Inner = A;
138    type Reader = MockReader;
139    type Writer = MockWriter;
140    type Lister = MockLister;
141    type Deleter = MockDeleter;
142
143    fn inner(&self) -> &Self::Inner {
144        &self.inner
145    }
146
147    async fn read(&self, path: &str, args: OpRead) -> Result<(RpRead, Self::Reader)> {
148        if let Some(reader_factory) = self.reader_factory.as_ref() {
149            let (rp_read, reader) = self.inner.read(path, args.clone()).await?;
150            let reader = reader_factory(path, args, Box::new(reader));
151            Ok((rp_read, MockReader { inner: reader }))
152        } else {
153            self.inner.read(path, args).await.map(|(rp_read, reader)| {
154                (
155                    rp_read,
156                    MockReader {
157                        inner: Box::new(reader),
158                    },
159                )
160            })
161        }
162    }
163
164    async fn write(&self, path: &str, args: OpWrite) -> Result<(RpWrite, Self::Writer)> {
165        if let Some(writer_factory) = self.writer_factory.as_ref() {
166            let (rp_write, writer) = self.inner.write(path, args.clone()).await?;
167            let writer = writer_factory(path, args, Box::new(writer));
168            Ok((rp_write, MockWriter { inner: writer }))
169        } else {
170            self.inner
171                .write(path, args)
172                .await
173                .map(|(rp_write, writer)| {
174                    (
175                        rp_write,
176                        MockWriter {
177                            inner: Box::new(writer),
178                        },
179                    )
180                })
181        }
182    }
183
184    async fn delete(&self) -> Result<(RpDelete, Self::Deleter)> {
185        if let Some(deleter_factory) = self.deleter_factory.as_ref() {
186            let (rp_delete, deleter) = self.inner.delete().await?;
187            let deleter = deleter_factory(Box::new(deleter));
188            Ok((rp_delete, MockDeleter { inner: deleter }))
189        } else {
190            self.inner.delete().await.map(|(rp_delete, deleter)| {
191                (
192                    rp_delete,
193                    MockDeleter {
194                        inner: Box::new(deleter),
195                    },
196                )
197            })
198        }
199    }
200
201    async fn list(&self, path: &str, args: OpList) -> Result<(RpList, Self::Lister)> {
202        if let Some(lister_factory) = self.lister_factory.as_ref() {
203            let (rp_list, lister) = self.inner.list(path, args.clone()).await?;
204            let lister = lister_factory(path, args, Box::new(lister));
205            Ok((rp_list, MockLister { inner: lister }))
206        } else {
207            self.inner.list(path, args).await.map(|(rp_list, lister)| {
208                (
209                    rp_list,
210                    MockLister {
211                        inner: Box::new(lister),
212                    },
213                )
214            })
215        }
216    }
217}
object_store/layers/mock.rs

object_store/layers/
mock.rs
