Import upstream v0.16.22, stripped
Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f Enterprise-only files removed or emptied: 63 Enterprise-only snippets removed: 117 in 50 files Dangling module declarations removed: 5 Cargo edits turning enterprise off: 14 Verification: clean Enterprise feature gates left for rebuilt features: 19 in 18 files Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.
This commit is contained in:
@@ -0,0 +1,300 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use crate::utils::{
|
||||
cleanup::{store_assert_is_empty, store_destroy},
|
||||
server::TestServerBuilder,
|
||||
};
|
||||
use registry::schema::structs::Rate;
|
||||
use registry::types::duration::Duration;
|
||||
use store::{InMemoryStore, dispatch::lookup::KeyValue};
|
||||
|
||||
#[tokio::test]
|
||||
pub async fn lookup_tests() {
|
||||
let test = TestServerBuilder::new("lookup_tests").await.build().await;
|
||||
let store = test.server.in_memory_store().clone();
|
||||
let rate = Rate {
|
||||
count: 1,
|
||||
period: Duration::from_millis(1000),
|
||||
};
|
||||
|
||||
println!(
|
||||
"Testing in-memory store {}...",
|
||||
std::env::var("MEMORY_STORE").unwrap_or_else(|_| "default".to_string())
|
||||
);
|
||||
if let InMemoryStore::Store(store) = &store {
|
||||
store_destroy(store).await;
|
||||
} else {
|
||||
// Reset redis counter
|
||||
store
|
||||
.key_set(KeyValue::new("abc", "0".as_bytes().to_vec()))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Test key
|
||||
let key = "xyz".as_bytes().to_vec();
|
||||
store
|
||||
.key_set(KeyValue::new(key.clone(), "world".to_string().into_bytes()))
|
||||
.await
|
||||
.unwrap();
|
||||
store.purge_in_memory_store().await.unwrap();
|
||||
assert_eq!(
|
||||
store.key_get::<String>(key.clone()).await.unwrap(),
|
||||
Some("world".to_string())
|
||||
);
|
||||
|
||||
// Test value expiry
|
||||
store
|
||||
.key_set(KeyValue::new(key.clone(), "hello".to_string().into_bytes()).expires(1))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store.key_get::<String>(key.clone()).await.unwrap(),
|
||||
Some("hello".to_string())
|
||||
);
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
|
||||
assert_eq!(None, store.key_get::<String>(key.clone()).await.unwrap());
|
||||
|
||||
store.purge_in_memory_store().await.unwrap();
|
||||
if let InMemoryStore::Store(store) = &store {
|
||||
store_assert_is_empty(store, store.clone().into(), false).await;
|
||||
}
|
||||
|
||||
// Test counter
|
||||
let key = "abc".as_bytes().to_vec();
|
||||
store
|
||||
.counter_incr(KeyValue::new(key.clone(), 1), true)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
|
||||
store
|
||||
.counter_incr(KeyValue::new(key.clone(), 2), true)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(3, store.counter_get(key.clone()).await.unwrap());
|
||||
store
|
||||
.counter_incr(KeyValue::new(key.clone(), -3), false)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
|
||||
|
||||
// Test counter expiry
|
||||
let key = "fgh".as_bytes().to_vec();
|
||||
store
|
||||
.counter_incr(KeyValue::new(key.clone(), 1).expires(1), false)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(1, store.counter_get(key.clone()).await.unwrap());
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
store.purge_in_memory_store().await.unwrap();
|
||||
assert_eq!(0, store.counter_get(key.clone()).await.unwrap());
|
||||
|
||||
// Test rate limiter
|
||||
assert!(
|
||||
store
|
||||
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
assert!(
|
||||
store
|
||||
.is_rate_allowed(0, "rate".as_bytes(), &rate, false)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
store.purge_in_memory_store().await.unwrap();
|
||||
if let InMemoryStore::Store(store) = &store {
|
||||
store_assert_is_empty(store, store.clone().into(), false).await;
|
||||
}
|
||||
|
||||
// Test locking
|
||||
for iteration in [1, 2] {
|
||||
let mut tasks = Vec::new();
|
||||
for _ in 0..100 {
|
||||
let store = store.clone();
|
||||
tasks.push(tokio::spawn(async move {
|
||||
store.try_lock(0, "lock".as_bytes(), 1).await.unwrap()
|
||||
}));
|
||||
}
|
||||
// Only one should return true
|
||||
let mut count = 0;
|
||||
for task in tasks {
|
||||
if task.await.unwrap() {
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
assert_eq!(1, count, "Iteration {}", iteration);
|
||||
|
||||
// Wait 2 seconds for the lock to expire
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
|
||||
}
|
||||
store.purge_in_memory_store().await.unwrap();
|
||||
|
||||
// Failed lock attempts must not extend the lock expiry
|
||||
assert!(store.try_lock(0, "lock".as_bytes(), 2).await.unwrap());
|
||||
let mut acquired = false;
|
||||
for _ in 0..12 {
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
||||
if store.try_lock(0, "lock".as_bytes(), 2).await.unwrap() {
|
||||
acquired = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(acquired, "Abandoned lock was never released");
|
||||
store.remove_lock(0, "lock".as_bytes()).await.unwrap();
|
||||
store.purge_in_memory_store().await.unwrap();
|
||||
if let InMemoryStore::Store(store) = &store {
|
||||
store_assert_is_empty(store, store.clone().into(), false).await;
|
||||
}
|
||||
|
||||
// Test prefix delete
|
||||
store
|
||||
.key_set(KeyValue::with_prefix(
|
||||
1,
|
||||
[0],
|
||||
"hello".to_string().into_bytes(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
for v in 0u32..2020u32 {
|
||||
store
|
||||
.key_set(KeyValue::with_prefix(
|
||||
0,
|
||||
pack_u32(0, v),
|
||||
"world".to_string().into_bytes(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.counter_incr(
|
||||
KeyValue::with_prefix(0, pack_u32(1, v), 123).expires(3600),
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Make sure the keys are there
|
||||
assert_eq!(
|
||||
Some("hello"),
|
||||
store
|
||||
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref()
|
||||
);
|
||||
for v in [0, 1000, 1001, 2000, 2001] {
|
||||
assert_eq!(
|
||||
Some("world"),
|
||||
store
|
||||
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref()
|
||||
);
|
||||
}
|
||||
for v in [0, 1000, 1001, 2000, 2001] {
|
||||
assert_ne!(
|
||||
0,
|
||||
store
|
||||
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
// Delete [0, 0, 0, 0, 1] prefix and make sure only the keys with that prefix are gone
|
||||
store
|
||||
.key_delete_prefix(&KeyValue::<()>::build_key(0, 1u32.to_be_bytes()))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Some("hello"),
|
||||
store
|
||||
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref()
|
||||
);
|
||||
for v in [0, 1000, 1001, 2000, 2001] {
|
||||
assert_eq!(
|
||||
Some("world"),
|
||||
store
|
||||
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref()
|
||||
);
|
||||
}
|
||||
|
||||
for v in [0, 1000, 1001, 2000, 2001] {
|
||||
assert_eq!(
|
||||
0,
|
||||
store
|
||||
.counter_get(KeyValue::<()>::build_key(0, pack_u32(1, v)))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
// Delete [0, 0, 0, 0, 0] prefix and make sure only the keys with that prefix are gone
|
||||
store
|
||||
.key_delete_prefix(&KeyValue::<()>::build_key(0, 0u32.to_be_bytes()))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
Some("hello"),
|
||||
store
|
||||
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref()
|
||||
);
|
||||
for v in [0, 1000, 1001, 2000, 2001] {
|
||||
assert_eq!(
|
||||
None,
|
||||
store
|
||||
.key_get::<String>(KeyValue::<()>::build_key(0, pack_u32(0, v)))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref()
|
||||
);
|
||||
}
|
||||
|
||||
// Delete [1, ...] prefix and make sure it's all gone
|
||||
store.key_delete_prefix(&[1u8]).await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
None,
|
||||
store
|
||||
.key_get::<String>(KeyValue::<()>::build_key(1, [0]))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref()
|
||||
);
|
||||
|
||||
if let InMemoryStore::Store(store) = &store {
|
||||
store_assert_is_empty(store, store.clone().into(), false).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn pack_u32(a: u32, b: u32) -> Vec<u8> {
|
||||
(((a as u64) << 32) | b as u64).to_be_bytes().to_vec()
|
||||
}
|
||||
Reference in New Issue
Block a user