Skip to main content

snix_castore/blobservice/
mocks.rs

1use std::{io, pin::pin, task::Poll};
2
3use tonic::async_trait;
4
5use crate::B3Digest;
6
7/// A BlobWriter simply accepting all data written to it.
8/// Used for testing purposes.
9pub struct TestBlobWriter {
10    buf: Vec<u8>,
11    closed: bool,
12}
13
14impl Default for TestBlobWriter {
15    fn default() -> Self {
16        Self::new()
17    }
18}
19
20impl TestBlobWriter {
21    pub fn new() -> Self {
22        Self {
23            buf: vec![],
24            closed: false,
25        }
26    }
27}
28
29impl tokio::io::AsyncWrite for TestBlobWriter {
30    fn poll_write(
31        self: std::pin::Pin<&mut Self>,
32        _cx: &mut std::task::Context<'_>,
33        buf: &[u8],
34    ) -> Poll<io::Result<usize>> {
35        pin!(self).buf.extend_from_slice(buf);
36        Poll::Ready(Ok(buf.len()))
37    }
38
39    fn poll_flush(
40        self: std::pin::Pin<&mut Self>,
41        _cx: &mut std::task::Context<'_>,
42    ) -> Poll<io::Result<()>> {
43        Poll::Ready(Ok(()))
44    }
45
46    fn poll_shutdown(
47        self: std::pin::Pin<&mut Self>,
48        _cx: &mut std::task::Context<'_>,
49    ) -> Poll<io::Result<()>> {
50        Poll::Ready(Ok(()))
51    }
52}
53
54#[async_trait]
55impl super::BlobWriter for TestBlobWriter {
56    async fn close(&mut self) -> io::Result<B3Digest> {
57        if self.closed {
58            Err(io::Error::other("already closed"))
59        } else {
60            Ok(blake3::hash(&self.buf).into())
61        }
62    }
63}