xmtp_db/
test_utils.rs

1#![allow(clippy::unwrap_used)]
2
3use std::path::Path;
4
5use crate::{DbConnection, EncryptedMessageStore, StorageOption};
6mod impls;
7mod mls_memory_storage;
8
9pub use mls_memory_storage::*;
10
11pub type TestDb = EncryptedMessageStore<crate::DefaultDatabase>;
12
13#[allow(async_fn_in_trait)]
14pub trait XmtpTestDb {
15    /// Create a validated, ephemeral database, running the migrations
16    async fn create_ephemeral_store() -> EncryptedMessageStore<crate::DefaultDatabase>;
17
18    async fn create_ephemeral_store_from_snapshot(
19        snapshot: &[u8],
20        path: Option<impl AsRef<Path>>,
21    ) -> EncryptedMessageStore<crate::DefaultDatabase>;
22
23    /// Create a validated, persistent database running the migrations
24    async fn create_persistent_store(
25        path: Option<String>,
26    ) -> EncryptedMessageStore<crate::DefaultDatabase>;
27    /// Create an empty database
28    /// does no validation and does not run migrations.
29    async fn create_database(path: Option<String>) -> crate::DefaultDatabase;
30}
31
32impl<Db> EncryptedMessageStore<Db> {
33    pub fn generate_enc_key() -> [u8; 32] {
34        let key = xmtp_common::rand_array::<32>();
35        tracing::debug!("generated key is [{}]", hex::encode(key));
36        key
37    }
38
39    #[cfg(not(target_arch = "wasm32"))]
40    pub fn remove_db_files<P: AsRef<str>>(path: P) {
41        use crate::database::native::EncryptedConnection;
42
43        let path = path.as_ref();
44        std::fs::remove_file(path).unwrap();
45        std::fs::remove_file(EncryptedConnection::salt_file(path).unwrap()).unwrap();
46    }
47
48    /// just a no-op on wasm32
49    #[cfg(target_arch = "wasm32")]
50    pub fn remove_db_files<P: AsRef<str>>(_path: P) {}
51}
52
53impl Clone for crate::MockXmtpDb {
54    fn clone(&self) -> Self {
55        panic!("clone is not allowed")
56    }
57}
58
59#[cfg(all(target_family = "wasm", target_os = "unknown"))]
60pub use wasm::*;
61#[cfg(all(target_family = "wasm", target_os = "unknown"))]
62mod wasm {
63    use super::*;
64    use crate::{PersistentOrMem, WasmDbConnection};
65    use futures::FutureExt;
66    use std::sync::Arc;
67
68    impl XmtpTestDb for super::TestDb {
69        async fn create_ephemeral_store() -> EncryptedMessageStore<crate::DefaultDatabase> {
70            let db = crate::database::WasmDb::new(&StorageOption::Ephemeral)
71                .await
72                .unwrap();
73            EncryptedMessageStore::new(db).unwrap()
74        }
75
76        async fn create_ephemeral_store_from_snapshot(
77            snapshot: &[u8],
78            _path: Option<impl AsRef<Path>>,
79        ) -> EncryptedMessageStore<crate::DefaultDatabase> {
80            let db = crate::database::WasmDb::new(&StorageOption::Ephemeral)
81                .await
82                .unwrap();
83            let store = EncryptedMessageStore::new_uninit(db).unwrap();
84            store
85                .db()
86                .raw_query_write(|conn| conn.deserialize_database_from_buffer(snapshot))
87                .unwrap();
88
89            store
90        }
91
92        async fn create_persistent_store(
93            path: Option<String>,
94        ) -> EncryptedMessageStore<crate::DefaultDatabase> {
95            let tmp = path.unwrap_or(xmtp_common::tmp_path());
96            let db = crate::database::WasmDb::new(&StorageOption::Persistent(tmp))
97                .await
98                .unwrap();
99            EncryptedMessageStore::new(db).unwrap()
100        }
101
102        async fn create_database(path: Option<String>) -> crate::DefaultDatabase {
103            let tmp = path.unwrap_or(xmtp_common::tmp_path());
104            crate::database::WasmDb::new(&StorageOption::Persistent(tmp))
105                .await
106                .unwrap()
107        }
108    }
109
110    /// Test harness that loads an Ephemeral store.
111    pub fn with_connection<F, R>(fun: F) -> R
112    where
113        F: FnOnce(
114            &crate::DbConnection<Arc<PersistentOrMem<WasmDbConnection, WasmDbConnection>>>,
115        ) -> R,
116    {
117        // ephemeral db connections do not use async so should resolve immediately
118        let db = crate::database::WasmDb::new(&StorageOption::Ephemeral)
119            .now_or_never()
120            .unwrap()
121            .unwrap();
122        let store = EncryptedMessageStore::new(db).unwrap();
123        let conn = store.conn();
124        fun(&DbConnection::new(conn))
125    }
126
127    /// Test harness that loads an Ephemeral store.
128    pub async fn with_connection_async<F, T, R>(fun: F) -> R
129    where
130        F: FnOnce(
131            crate::DbConnection<Arc<PersistentOrMem<WasmDbConnection, WasmDbConnection>>>,
132        ) -> T,
133        T: Future<Output = R>,
134    {
135        let db = crate::database::WasmDb::new(&StorageOption::Ephemeral)
136            .await
137            .unwrap();
138        let store = EncryptedMessageStore::new(db).unwrap();
139        let conn = store.conn();
140        fun(DbConnection::new(conn)).await
141    }
142
143    impl EncryptedMessageStore<crate::database::WasmDb> {
144        pub async fn new_test() -> Self {
145            let db = crate::database::WasmDb::new(&StorageOption::Ephemeral)
146                .await
147                .unwrap();
148            EncryptedMessageStore::new(db).expect("constructing message store failed.")
149        }
150
151        pub async fn new_test_with_path(path: &str) -> Self {
152            let db = crate::database::WasmDb::new(&StorageOption::Persistent(path.into()))
153                .await
154                .unwrap();
155            EncryptedMessageStore::new(db).expect("constructing message store failed.")
156        }
157    }
158}
159
160#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
161pub use native::*;
162#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
163mod native {
164    use super::*;
165    use crate::{
166        ConnectionExt, EphemeralDbConnection, MIGRATIONS, NativeDbConnection, PersistentOrMem,
167    };
168    use diesel::{Connection, SqliteConnection, connection::SimpleConnection};
169    use diesel_migrations::MigrationHarness;
170    use std::sync::Arc;
171
172    impl XmtpTestDb for super::TestDb {
173        async fn create_ephemeral_store() -> crate::DefaultStore {
174            let opts = StorageOption::Ephemeral;
175            let db = crate::database::NativeDb::new_unencrypted(&opts).unwrap();
176            EncryptedMessageStore::new(db).unwrap()
177        }
178        async fn create_ephemeral_store_from_snapshot(
179            mut snapshot: &[u8],
180            path: Option<impl AsRef<Path>>,
181        ) -> crate::DefaultStore {
182            let path = path.as_ref();
183            let mut buffer;
184
185            let mut i = 0;
186            let store = loop {
187                let opts = StorageOption::Ephemeral;
188                let db = crate::database::NativeDb::new_unencrypted(&opts).unwrap();
189                let store = EncryptedMessageStore::new_uninit(db).unwrap();
190                let result = store.db().raw_query_write(|conn| {
191                    conn.deserialize_database_from_buffer(snapshot)?;
192                    conn.batch_execute("PRAGMA journal_mode = DELETE")?;
193                    Ok(())
194                });
195
196                if result.is_ok() {
197                    break store;
198                } else if i >= 1 || path.is_none() {
199                    #[allow(clippy::panicking_unwrap)]
200                    result.unwrap();
201                }
202
203                if let Some(path) = path {
204                    let path = path.as_ref();
205                    // WAL is not compatible with ephemeral databases. Attempt to update and try one more time.
206                    {
207                        let mut conn =
208                            SqliteConnection::establish(path.to_string_lossy().as_ref()).unwrap();
209                        conn.batch_execute("PRAGMA journal_mode = DELETE").unwrap();
210                    }
211
212                    buffer = std::fs::read(path).unwrap();
213                    snapshot = &buffer;
214                };
215
216                i += 1;
217            };
218
219            store
220                .conn()
221                .raw_query_write(|c| {
222                    c.run_pending_migrations(MIGRATIONS).unwrap();
223                    Ok(())
224                })
225                .unwrap();
226
227            store
228        }
229        async fn create_persistent_store(path: Option<String>) -> crate::DefaultStore {
230            let path = path.unwrap_or(xmtp_common::tmp_path());
231            let opts = StorageOption::Persistent(path.to_string());
232            let db = crate::database::NativeDb::new(&opts, [0u8; 32]).unwrap();
233            EncryptedMessageStore::new(db).expect("constructing message store failed.")
234        }
235
236        async fn create_database(path: Option<String>) -> crate::DefaultDatabase {
237            let path = path.unwrap_or(xmtp_common::tmp_path());
238            let opts = StorageOption::Persistent(path.to_string());
239            crate::database::NativeDb::new(&opts, xmtp_common::rand_array::<32>()).unwrap()
240        }
241    }
242
243    /// Test harness that loads an Ephemeral store.
244    pub fn with_connection<F, R>(fun: F) -> R
245    where
246        F: FnOnce(
247            &crate::DbConnection<Arc<PersistentOrMem<NativeDbConnection, EphemeralDbConnection>>>,
248        ) -> R,
249    {
250        let opts = StorageOption::Ephemeral;
251        let db = crate::database::NativeDb::new_unencrypted(&opts).unwrap();
252        let store = EncryptedMessageStore::new(db).unwrap();
253        let conn = store.conn();
254        fun(&DbConnection::new(conn))
255    }
256
257    /// Test harness that loads an Ephemeral store.
258    pub async fn with_connection_async<F, T, R>(fun: F) -> R
259    where
260        F: FnOnce(
261            crate::DbConnection<Arc<PersistentOrMem<NativeDbConnection, EphemeralDbConnection>>>,
262        ) -> T,
263        T: Future<Output = R>,
264    {
265        let opts = StorageOption::Ephemeral;
266        let db = crate::database::NativeDb::new_unencrypted(&opts).unwrap();
267        let store = EncryptedMessageStore::new(db).unwrap();
268        let conn = store.conn();
269        fun(DbConnection::new(conn)).await
270    }
271
272    impl EncryptedMessageStore<crate::database::NativeDb> {
273        pub async fn new_test() -> Self {
274            let tmp_path = xmtp_common::tmp_path();
275            let opts = StorageOption::Persistent(tmp_path);
276            let db =
277                crate::database::NativeDb::new(&opts, xmtp_common::rand_array::<32>()).unwrap();
278            EncryptedMessageStore::new(db).expect("constructing message store failed.")
279        }
280
281        pub async fn new_test_with_path(path: &str) -> Self {
282            let opts = StorageOption::Persistent(path.to_string());
283            let db =
284                crate::database::NativeDb::new(&opts, xmtp_common::rand_array::<32>()).unwrap();
285            EncryptedMessageStore::new(db).expect("constructing message store failed.")
286        }
287    }
288}
xmtp_db/test_utils.rs

xmtp_db/
test_utils.rs
