From d7ae022005312465c4800dd455ed8c6c03fafa65 Mon Sep 17 00:00:00 2001 From: Alex Snaps Date: Wed, 13 Jul 2022 15:00:49 +0300 Subject: [PATCH] We need to insert a HashSet before doing anything --- limitador/src/storage/in_memory.rs | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/limitador/src/storage/in_memory.rs b/limitador/src/storage/in_memory.rs index 396441de..64db71ed 100644 --- a/limitador/src/storage/in_memory.rs +++ b/limitador/src/storage/in_memory.rs @@ -132,8 +132,8 @@ impl InMemoryStorage { { Entry::Occupied(mut e) => { e.get_mut() - .get_mut(counter.limit()) - .unwrap() + .entry(counter.limit().clone()) + .or_default() .insert(counter.clone()); } Entry::Vacant(e) => { @@ -181,3 +181,25 @@ impl Default for InMemoryStorage { Self::new(1000) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn counters_for_multiple_limit_per_ns() { + let storage = InMemoryStorage::new(100); + let namespace = "test_namespace"; + let limit_1 = Limit::new(namespace, 1, 1, vec!["req.method == GET"], vec!["app_id"]); + let limit_2 = Limit::new(namespace, 1, 10, vec!["req.method == GET"], vec!["app_id"]); + let counter_1 = Counter::new(limit_1, HashMap::default()); + let counter_2 = Counter::new(limit_2, HashMap::default()); + storage.update_counter(&counter_1, 1).unwrap(); + storage.update_counter(&counter_2, 1).unwrap(); + + assert_eq!( + storage.counters_in_namespace(counter_1.namespace()).len(), + 2 + ); + } +}