1
mod compatibility;
2

            
3
use std::time::Duration;
4

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

            
16
use crate::{
17
    config::{Builder, StorageConfiguration},
18
    Database, Storage,
19
};
20

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

            
29
                use super::*;
30
                use crate::{AsyncDatabase, AsyncStorage};
31
                struct AsyncTestHarness {
32
                    _directory: TestDirectory,
33
                    db: AsyncDatabase,
34
                    storage: AsyncStorage,
35
                }
36

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

            
47
                        #[cfg(feature = "compression")]
48
                        {
49
                            config = config.default_compression(crate::config::Compression::Lz4);
50
                        }
51

            
52
                        let storage = AsyncStorage::open(config).await?;
53
                        let db = storage
54
                            .create_database::<BasicSchema>("tests", false)
55
                            .await?;
56

            
57
                        Ok(Self {
58
                            _directory: directory,
59
                            storage,
60
                            db,
61
                        })
62
                    }
63

            
64
4
                    const fn server_name() -> &'static str {
65
4
                        stringify!($name)
66
4
                    }
67

            
68
4
                    fn server(&self) -> &'_ AsyncStorage {
69
4
                        &self.storage
70
4
                    }
71

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

            
84
60
                    async fn connect(&self) -> anyhow::Result<AsyncDatabase> {
85
60
                        Ok(self.db.clone())
86
60
                    }
87

            
88
52
                    pub async fn shutdown(&self) -> anyhow::Result<()> {
89
52
                        Ok(())
90
52
                    }
91
                }
92

            
93
                bonsaidb_core::define_async_connection_test_suite!(AsyncTestHarness);
94

            
95
                bonsaidb_core::define_async_pubsub_test_suite!(AsyncTestHarness);
96

            
97
                bonsaidb_core::define_async_kv_test_suite!(AsyncTestHarness);
98
            }
99
            mod blocking {
100
                use bonsaidb_core::connection::StorageConnection;
101

            
102
                use super::*;
103
                struct BlockingTestHarness {
104
                    _directory: TestDirectory,
105
                    db: Database,
106
                    storage: Storage,
107
                }
108

            
109
                impl BlockingTestHarness {
110
62
                    fn new(test: HarnessTest) -> anyhow::Result<Self> {
111
62
                        let directory =
112
62
                            TestDirectory::new(format!("blocking-{}-{}", stringify!($name), test));
113
62
                        let mut config =
114
62
                            StorageConfiguration::new(&directory).with_schema::<BasicSchema>()?;
115
62
                        if stringify!($name) == "memory" {
116
31
                            config = config.memory_only()
117
31
                        }
118

            
119
                        #[cfg(feature = "compression")]
120
62
                        {
121
62
                            config = config.default_compression(crate::config::Compression::Lz4);
122
62
                        }
123

            
124
62
                        let storage = Storage::open(config)?;
125
62
                        let db = storage.create_database::<BasicSchema>("tests", false)?;
126

            
127
62
                        Ok(Self {
128
62
                            _directory: directory,
129
62
                            storage,
130
62
                            db,
131
62
                        })
132
62
                    }
133

            
134
4
                    const fn server_name() -> &'static str {
135
4
                        stringify!($name)
136
4
                    }
137

            
138
4
                    fn server(&self) -> &'_ Storage {
139
4
                        &self.storage
140
4
                    }
141

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

            
154
60
                    fn connect(&self) -> anyhow::Result<Database> {
155
60
                        Ok(self.db.clone())
156
60
                    }
157

            
158
52
                    pub fn shutdown(&self) -> anyhow::Result<()> {
159
52
                        Ok(())
160
52
                    }
161
                }
162

            
163
                bonsaidb_core::define_blocking_connection_test_suite!(BlockingTestHarness);
164

            
165
                bonsaidb_core::define_blocking_pubsub_test_suite!(BlockingTestHarness);
166

            
167
                bonsaidb_core::define_blocking_kv_test_suite!(BlockingTestHarness);
168
            }
169
        }
170
    };
171
}
172

            
173
400
define_local_suite!(persisted);
174
400
define_local_suite!(memory);
175

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

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

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

            
208
        // Regular query should show the correct data
209
1
        assert_eq!(db.view::<BasicByBrokenParentId>().query()?.len(), 1);
210
    }
211

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

            
228
    unreachable!("Integrity checker didn't run in the allocated time")
229
1
}
230

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

            
239
1
        let document_header = db
240
1
            .collection::<EncryptedBasic>()
241
1
            .push(&EncryptedBasic::new("hello"))?;
242

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

            
250
1
        document_header
251
1
    };
252
1

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

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

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

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

            
270
1
    Ok(())
271
1
}
272

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

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

            
294
1
            db.set_key("a", &0_u32)
295
1
                .expire_in(Duration::from_secs(3))
296
1
                .execute()?;
297
        }
298

            
299
        {
300
1
            let db = Database::open::<()>(StorageConfiguration::new(&path))?;
301

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

            
315
1
            timing.wait_until(Duration::from_secs(4));
316

            
317
1
            assert!(db.get_key("a").query()?.is_none());
318
        }
319

            
320
1
        break;
321
1
    }
322
1
    Ok(())
323
1
}