diff --git a/frame/support/src/storage/types/counted_map.rs b/frame/support/src/storage/types/counted_map.rs index 8c19434767f49..3361c4093dce0 100644 --- a/frame/support/src/storage/types/counted_map.rs +++ b/frame/support/src/storage/types/counted_map.rs @@ -134,7 +134,10 @@ where /// Store or remove the value to be associated with `key` so that `get` returns the `query`. pub fn set>(key: KeyArg, q: QueryKind::Query) { - ::Map::set(key, q) + match QueryKind::from_query_to_optional_value(q) { + Some(v) => Self::insert(key, v), + None => Self::remove(key), + } } /// Swap the values of two keys. @@ -745,6 +748,22 @@ mod test { // Test initialize_counter. assert_eq!(A::initialize_counter(), 2); + + // Set non-existing. + A::set(30, 100); + + assert_eq!(A::contains_key(30), true); + assert_eq!(A::get(30), 100); + assert_eq!(A::try_get(30), Ok(100)); + assert_eq!(A::count(), 3); + + // Set existing. + A::set(30, 101); + + assert_eq!(A::contains_key(30), true); + assert_eq!(A::get(30), 101); + assert_eq!(A::try_get(30), Ok(101)); + assert_eq!(A::count(), 3); }) } @@ -976,6 +995,40 @@ mod test { // Test initialize_counter. assert_eq!(B::initialize_counter(), 2); + + // Set non-existing. + B::set(30, Some(100)); + + assert_eq!(B::contains_key(30), true); + assert_eq!(B::get(30), Some(100)); + assert_eq!(B::try_get(30), Ok(100)); + assert_eq!(B::count(), 3); + + // Set existing. + B::set(30, Some(101)); + + assert_eq!(B::contains_key(30), true); + assert_eq!(B::get(30), Some(101)); + assert_eq!(B::try_get(30), Ok(101)); + assert_eq!(B::count(), 3); + + // Unset existing. + B::set(30, None); + + assert_eq!(B::contains_key(30), false); + assert_eq!(B::get(30), None); + assert_eq!(B::try_get(30), Err(())); + + assert_eq!(B::count(), 2); + + // Unset non-existing. + B::set(31, None); + + assert_eq!(B::contains_key(31), false); + assert_eq!(B::get(31), None); + assert_eq!(B::try_get(31), Err(())); + + assert_eq!(B::count(), 2); }) } diff --git a/frame/support/src/storage/types/map.rs b/frame/support/src/storage/types/map.rs index 0f89e2378a55d..3e2b744d72f89 100644 --- a/frame/support/src/storage/types/map.rs +++ b/frame/support/src/storage/types/map.rs @@ -627,6 +627,48 @@ mod test { assert_eq!(AValueQueryWithAnOnEmpty::take(2), 97); assert_eq!(A::contains_key(2), false); + // Set non-existing. + B::set(30, 100); + + assert_eq!(B::contains_key(30), true); + assert_eq!(B::get(30), 100); + assert_eq!(B::try_get(30), Ok(100)); + + // Set existing. + B::set(30, 101); + + assert_eq!(B::contains_key(30), true); + assert_eq!(B::get(30), 101); + assert_eq!(B::try_get(30), Ok(101)); + + // Set non-existing. + A::set(30, Some(100)); + + assert_eq!(A::contains_key(30), true); + assert_eq!(A::get(30), Some(100)); + assert_eq!(A::try_get(30), Ok(100)); + + // Set existing. + A::set(30, Some(101)); + + assert_eq!(A::contains_key(30), true); + assert_eq!(A::get(30), Some(101)); + assert_eq!(A::try_get(30), Ok(101)); + + // Unset existing. + A::set(30, None); + + assert_eq!(A::contains_key(30), false); + assert_eq!(A::get(30), None); + assert_eq!(A::try_get(30), Err(())); + + // Unset non-existing. + A::set(31, None); + + assert_eq!(A::contains_key(31), false); + assert_eq!(A::get(31), None); + assert_eq!(A::try_get(31), Err(())); + B::insert(2, 10); assert_eq!(A::migrate_key::(2), Some(10)); assert_eq!(A::contains_key(2), true);