1
mod compatibility;
2

            
3
use std::time::Duration;
4

            
5
use bonsaidb_core::connection::{AccessPolicy, Connection};
6
use bonsaidb_core::permissions::{Permissions, Statement};
7
#[cfg(feature = "encryption")]
8
use bonsaidb_core::test_util::EncryptedBasic;
9
use bonsaidb_core::test_util::{
10
    Basic, BasicByBrokenParentId, BasicByParentId, BasicCollectionWithNoViews,
11
    BasicCollectionWithOnlyBrokenParentId, BasicSchema, HarnessTest, TestDirectory,
12
};
13

            
14
use crate::config::{Builder, StorageConfiguration};
15
use crate::{Database, Storage};
16

            
17
macro_rules! define_local_suite {
18
    ($name:ident) => {
19
        mod $name {
20
            use super::*;
21
            #[cfg(feature = "async")]
22
            mod r#async {
23
                use bonsaidb_core::connection::AsyncStorageConnection;
24

            
25
                use super::*;
26
                use crate::{AsyncDatabase, AsyncStorage};
27
                struct AsyncTestHarness {
28
                    _directory: TestDirectory,
29
                    db: AsyncDatabase,
30
                    storage: AsyncStorage,
31
                }
32

            
33
                impl AsyncTestHarness {
34
70
                    async fn new(test: HarnessTest) -> anyhow::Result<Self> {
35
                        let directory =
36
                            TestDirectory::new(format!("async-{}-{}", stringify!($name), test));
37
                        let mut config =
38
                            StorageConfiguration::new(&directory).with_schema::<BasicSchema>()?;
39
                        if stringify!($name) == "memory" {
40
                            config = config.memory_only()
41
                        }
42

            
43
                        #[cfg(feature = "compression")]
44
                        {
45
                            config = config.default_compression(crate::config::Compression::Lz4);
46
                        }
47

            
48
                        let storage = AsyncStorage::open(config).await?;
49
                        let db = storage
50
                            .create_database::<BasicSchema>("tests", false)
51
                            .await?;
52

            
53
                        Ok(Self {
54
                            _directory: directory,
55
                            storage,
56
                            db,
57
                        })
58
                    }
59

            
60
6
                    const fn server_name() -> &'static str {
61
6
                        stringify!($name)
62
6
                    }
63

            
64
6
                    fn server(&self) -> &'_ AsyncStorage {
65
6
                        &self.storage
66
6
                    }
67

            
68
                    #[allow(dead_code)]
69
                    async fn connect_with_permissions(
70
                        &self,
71
                        permissions: Vec<Statement>,
72
                        _label: &str,
73
                    ) -> anyhow::Result<AsyncDatabase> {
74
                        Ok(self
75
                            .db
76
                            .with_effective_permissions(Permissions::from(permissions))
77
                            .unwrap())
78
                    }
79

            
80
68
                    async fn connect(&self) -> anyhow::Result<AsyncDatabase> {
81
68
                        Ok(self.db.clone())
82
68
                    }
83

            
84
60
                    pub async fn shutdown(&self) -> anyhow::Result<()> {
85
60
                        Ok(())
86
60
                    }
87
                }
88

            
89
                bonsaidb_core::define_async_connection_test_suite!(AsyncTestHarness);
90

            
91
                bonsaidb_core::define_async_pubsub_test_suite!(AsyncTestHarness);
92

            
93
                bonsaidb_core::define_async_kv_test_suite!(AsyncTestHarness);
94
            }
95
            mod blocking {
96
                use bonsaidb_core::connection::StorageConnection;
97

            
98
                use super::*;
99
                struct BlockingTestHarness {
100
                    _directory: TestDirectory,
101
                    db: Database,
102
                    storage: Storage,
103
                }
104

            
105
                impl BlockingTestHarness {
106
72
                    fn new(test: HarnessTest) -> anyhow::Result<Self> {
107
72
                        let directory =
108
72
                            TestDirectory::new(format!("blocking-{}-{}", stringify!($name), test));
109
72
                        let mut config =
110
72
                            StorageConfiguration::new(&directory).with_schema::<BasicSchema>()?;
111
72
                        if stringify!($name) == "memory" {
112
36
                            config = config.memory_only()
113
36
                        }
114

            
115
                        #[cfg(feature = "compression")]
116
72
                        {
117
72
                            config = config.default_compression(crate::config::Compression::Lz4);
118
72
                        }
119

            
120
72
                        let storage = Storage::open(config)?;
121
72
                        let db = storage.create_database::<BasicSchema>("tests", false)?;
122

            
123
72
                        Ok(Self {
124
72
                            _directory: directory,
125
72
                            storage,
126
72
                            db,
127
72
                        })
128
72
                    }
129

            
130
6
                    const fn server_name() -> &'static str {
131
6
                        stringify!($name)
132
6
                    }
133

            
134
6
                    fn server(&self) -> &'_ Storage {
135
6
                        &self.storage
136
6
                    }
137

            
138
                    #[allow(dead_code)]
139
                    fn connect_with_permissions(
140
                        &self,
141
                        permissions: Vec<Statement>,
142
                        _label: &str,
143
                    ) -> anyhow::Result<Database> {
144
                        Ok(self
145
                            .db
146
                            .with_effective_permissions(Permissions::from(permissions))
147
                            .unwrap())
148
                    }
149

            
150
70
                    fn connect(&self) -> anyhow::Result<Database> {
151
70
                        Ok(self.db.clone())
152
70
                    }
153

            
154
62
                    pub fn shutdown(&self) -> anyhow::Result<()> {
155
62
                        Ok(())
156
62
                    }
157
                }
158

            
159
                bonsaidb_core::define_blocking_connection_test_suite!(BlockingTestHarness);
160

            
161
                bonsaidb_core::define_blocking_pubsub_test_suite!(BlockingTestHarness);
162

            
163
                bonsaidb_core::define_blocking_kv_test_suite!(BlockingTestHarness);
164
            }
165
        }
166
    };
167
}
168

            
169
400
define_local_suite!(persisted);
170
400
define_local_suite!(memory);
171

            
172
1
#[test]
173
#[cfg_attr(not(feature = "compression"), allow(unused_mut))]
174
1
fn integrity_checks() -> anyhow::Result<()> {
175
1
    let path = TestDirectory::new("integrity-checks");
176
1
    let mut config = StorageConfiguration::new(&path);
177
1
    #[cfg(feature = "compression")]
178
1
    {
179
1
        config = config.default_compression(crate::config::Compression::Lz4);
180
1
    }
181
    // To ensure full cleanup between each block, each runs in its own runtime;
182

            
183
    // Add a doc with no views installed
184
    {
185
1
        let db = Database::open::<BasicCollectionWithNoViews>(config.clone())?;
186
1
        let collection = db.collection::<BasicCollectionWithNoViews>();
187
1
        collection.push(&Basic::default().with_parent_id(1))?;
188
    }
189
    // Connect with a new view and see the automatic update with a query
190
    {
191
1
        let db = Database::open::<BasicCollectionWithOnlyBrokenParentId>(config.clone())?;
192
        // Give the integrity scanner time to run if it were to run (it shouldn't in this configuration).
193
1
        std::thread::sleep(Duration::from_millis(100));
194

            
195
        // NoUpdate should return data without the validation checker having run.
196
1
        assert_eq!(
197
1
            db.view::<BasicByBrokenParentId>()
198
1
                .with_access_policy(AccessPolicy::NoUpdate)
199
1
                .query()?
200
1
                .len(),
201
            0
202
        );
203

            
204
        // Regular query should show the correct data
205
1
        assert_eq!(db.view::<BasicByBrokenParentId>().query()?.len(), 1);
206
    }
207

            
208
    // Connect with a fixed view, and wait for the integrity scanner to work
209
1
    let db = Database::open::<Basic>(config.check_view_integrity_on_open(true))?;
210
1
    for _ in 0_u8..100 {
211
1
        std::thread::sleep(Duration::from_millis(1000));
212
1
        if db
213
1
            .view::<BasicByParentId>()
214
1
            .with_access_policy(AccessPolicy::NoUpdate)
215
1
            .with_key(&Some(1))
216
1
            .query()?
217
1
            .len()
218
            == 1
219
        {
220
1
            return Ok(());
221
        }
222
    }
223

            
224
    unreachable!("Integrity checker didn't run in the allocated time")
225
1
}
226

            
227
1
#[test]
228
#[cfg(feature = "encryption")]
229
1
fn encryption() -> anyhow::Result<()> {
230
1
    use bonsaidb_core::schema::SerializedCollection;
231
1
    let path = TestDirectory::new("encryption");
232
1
    let document_header = {
233
1
        let db = Database::open::<BasicSchema>(StorageConfiguration::new(&path))?;
234

            
235
1
        let document_header = db
236
1
            .collection::<EncryptedBasic>()
237
1
            .push(&EncryptedBasic::new("hello"))?;
238

            
239
        // Retrieve the document, showing that it was stored successfully.
240
1
        let doc = db
241
1
            .collection::<EncryptedBasic>()
242
1
            .get(&document_header.id)?
243
1
            .expect("doc not found");
244
1
        assert_eq!(&EncryptedBasic::document_contents(&doc)?.value, "hello");
245

            
246
1
        document_header
247
1
    };
248
1

            
249
1
    // By resetting the encryption key, we should be able to force an error in
250
1
    // decryption, which proves that the document was encrypted. To ensure the
251
1
    // server starts up and generates a new key, we must delete the sealing key.
252
1

            
253
1
    std::fs::remove_file(path.join("master-keys"))?;
254

            
255
1
    let db = Database::open::<BasicSchema>(StorageConfiguration::new(&path))?;
256

            
257
    // Try retrieving the document, but expect an error decrypting.
258
1
    if let Err(bonsaidb_core::Error::Other { error, .. }) =
259
1
        db.collection::<EncryptedBasic>().get(&document_header.id)
260
    {
261
1
        assert!(error.contains("vault"));
262
    } else {
263
        panic!("successfully retrieved encrypted document without keys");
264
    }
265

            
266
1
    Ok(())
267
1
}
268

            
269
1
#[test]
270
1
fn expiration_after_close() -> anyhow::Result<()> {
271
    use bonsaidb_core::keyvalue::KeyValue;
272
    use bonsaidb_core::test_util::TimingTest;
273
1
    loop {
274
1
        let path = TestDirectory::new("expiration-after-close");
275
1
        // To ensure full cleanup between each block, each runs in its own runtime;
276
1
        let timing = TimingTest::new(Duration::from_millis(100));
277
        // Set a key with an expiration, then close it. Then try to validate it
278
        // exists after opening, and then expires at the correct time.
279
        {
280
1
            let db = Database::open::<()>(StorageConfiguration::new(&path))?;
281

            
282
            // TODO This is a workaroun for the key-value expiration task
283
            // taking ownership of an instance of Database. If this async
284
            // task runs too quickly, sometimes things don't get cleaned up
285
            // if that task hasn't completed. This pause ensures the startup
286
            // tasks complete before we continue with the test. This should
287
            // be replaced with a proper shutdown call for the local
288
            // storage/database.
289
1
            std::thread::sleep(Duration::from_millis(100));
290
1

            
291
1
            db.set_key("a", &0_u32)
292
1
                .expire_in(Duration::from_secs(3))
293
1
                .execute()?;
294
        }
295

            
296
        {
297
1
            let db = Database::open::<()>(StorageConfiguration::new(&path))?;
298

            
299
1
            let key = db.get_key("a").query()?;
300
            // Due to not having a reliable way to shut down the database,
301
            // we can't make many guarantees about what happened after
302
            // setting the key in the above block. If we get None back,
303
            // we'll consider the test needing to retry. Once we have a
304
            // shutdown operation that guarantees that the key-value store
305
            // persists, the key.is_none() check shoud be removed, instead
306
            // asserting `key.is_some()`.
307
1
            if timing.elapsed() > Duration::from_secs(1) || key.is_none() {
308
                println!("Retrying  expiration_after_close because it was too slow");
309
                continue;
310
1
            }
311
1

            
312
1
            timing.wait_until(Duration::from_secs(4));
313

            
314
1
            assert!(db.get_key("a").query()?.is_none());
315
        }
316

            
317
1
        break;
318
1
    }
319
1
    Ok(())
320
1
}