snix_castore/blobstore/
mock_writer.rs

1//! A [super::Writer] that can be used for tests
2
3use tonic::async_trait;
4
5use crate::{B3Digest, blobstore::Writer};
6
7/// Implements [super::Writer] and panics if the expected data wasn't written on close.
8pub struct MockWriter {
9    hasher: parking_lot::RwLock<blake3::Hasher>,
10    expected_digest: B3Digest,
11    closed: bool,
12}
13
14impl MockWriter {
15    /// Constructs a new [MockWriter] expecting the passed digest.
16    pub fn new(expected_digest: B3Digest) -> Self {
17        Self {
18            hasher: Default::default(),
19            expected_digest,
20            closed: false,
21        }
22    }
23}
24
25impl Drop for MockWriter {
26    fn drop(&mut self) {
27        // Only panic if the test isn't already panicking for another reason (as in test cases for MockWriter itself)
28        if !self.closed && !std::thread::panicking() {
29            panic!(
30                "MockWriter with expected digest {} was not closed",
31                self.expected_digest
32            );
33        }
34    }
35}
36
37impl tokio::io::AsyncWrite for MockWriter {
38    fn poll_write(
39        self: std::pin::Pin<&mut Self>,
40        _cx: &mut std::task::Context<'_>,
41        buf: &[u8],
42    ) -> std::task::Poll<std::io::Result<usize>> {
43        if self.closed {
44            return std::task::Poll::Ready(Err(std::io::Error::new(
45                std::io::ErrorKind::BrokenPipe,
46                "already closed",
47            )));
48        }
49        self.hasher.write().update(buf);
50        std::task::Poll::Ready(Ok(buf.len()))
51    }
52
53    fn poll_flush(
54        self: std::pin::Pin<&mut Self>,
55        _cx: &mut std::task::Context<'_>,
56    ) -> std::task::Poll<std::io::Result<()>> {
57        std::task::Poll::Ready(Ok(()))
58    }
59
60    fn poll_shutdown(
61        self: std::pin::Pin<&mut Self>,
62        _cx: &mut std::task::Context<'_>,
63    ) -> std::task::Poll<std::io::Result<()>> {
64        std::task::Poll::Ready(Ok(()))
65    }
66}
67
68#[async_trait]
69impl Writer for MockWriter {
70    async fn close(&mut self) -> std::io::Result<crate::B3Digest> {
71        let actual_digest: B3Digest = self.hasher.read().finalize().into();
72
73        if self.expected_digest != actual_digest {
74            panic!(
75                "unexpected digest {} encountered, expected {}",
76                actual_digest, self.expected_digest
77            )
78        }
79
80        self.closed = true;
81
82        Ok(actual_digest)
83    }
84}
85
86#[cfg(test)]
87mod test {
88    use tokio::io::AsyncWriteExt;
89
90    use super::MockWriter;
91    use crate::{blobstore::Writer, chunkstore::Chunk};
92
93    #[tokio::test]
94    async fn ok() {
95        let digest = Chunk::from_static(b"foo").digest();
96        let mut mock_writer = MockWriter::new(digest);
97        mock_writer
98            .write_all(b"foo")
99            .await
100            .expect("write to succeed");
101
102        let actual_digest = mock_writer.close().await.expect("close to succeed");
103        assert_eq!(
104            digest, actual_digest,
105            "expect returned digest to match data written into"
106        );
107    }
108
109    #[tokio::test]
110    #[should_panic]
111    async fn panic_wrong_data_written() {
112        let digest = Chunk::from_static(b"foo").digest();
113        let mut mock_writer = MockWriter::new(digest);
114        mock_writer
115            .write_all(b"bar")
116            .await
117            .expect("write to succeed");
118
119        // panics
120        let _ = mock_writer.close().await;
121    }
122
123    #[tokio::test]
124    #[should_panic]
125    async fn panic_drop_without_close() {
126        let digest = Chunk::from_static(b"foo").digest();
127        let mut mock_writer = MockWriter::new(digest);
128        mock_writer
129            .write_all(b"foo")
130            .await
131            .expect("write to succeed");
132
133        // panics on drop
134        drop(mock_writer)
135    }
136}
snix_castore/blobstore/mock_writer.rs

snix_castore/blobstore/
mock_writer.rs
