Prefer usages of containsOnlyKeys when asserting keySets of a map

PicnicSupermarket · Apr 16, 2023 · 1fbff8c · 1fbff8c
1 parent 929f1dd
commit 1fbff8c
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 0 deletions.
diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/AssertJMapRules.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/AssertJMapRules.java
@@ -21,6 +21,7 @@
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.TreeMap;
 import java.util.TreeSet;
 import org.assertj.core.api.AbstractAssert;
@@ -220,6 +221,19 @@ MapAssert<K, V> after(Map<K, V> map, K key) {
     }
   }
 
+  static final class AssertThatMapContainsOnlyKeys<K, V> {
+    @BeforeTemplate
+    AbstractAssert<?, ?> before(Map<K, V> map, Set<K> keys) {
+      return assertThat(map.keySet()).hasSameElementsAs(keys);
+    }
+
+    @AfterTemplate
+    @UseImportPolicy(STATIC_IMPORT_ALWAYS)
+    AbstractAssert<?, ?> after(Map<K, V> map, Set<K> keys) {
+      return assertThat(map).containsOnlyKeys(keys);
+    }
+  }
+
   static final class AssertThatMapContainsValue<K, V> {
     @BeforeTemplate
     AbstractBooleanAssert<?> before(Map<K, V> map, V value) {

diff --git a/...rib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java b/...rib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java
@@ -130,6 +130,10 @@ MapAssert<Integer, Integer> testAbstractMapAssertHasSameSizeAs() {
     return assertThat(ImmutableMap.of(1, 2).containsKey(3)).isFalse();
   }
 
+  AbstractAssert<?, ?> testAssertThatMapContainsOnlyKeys() {
+    return assertThat(ImmutableMap.of(1, 2).keySet()).hasSameElementsAs(ImmutableSet.of(1));
+  }
+
   AbstractAssert<?, ?> testAssertThatMapContainsValue() {
     return assertThat(ImmutableMap.of(1, 2).containsValue(3)).isTrue();
   }

diff --git a/...ib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java b/...ib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java
@@ -130,6 +130,10 @@ MapAssert<Integer, Integer> testAbstractMapAssertHasSameSizeAs() {
     return assertThat(ImmutableMap.of(1, 2)).doesNotContainKey(3);
   }
 
+  AbstractAssert<?, ?> testAssertThatMapContainsOnlyKeys() {
+    return assertThat(ImmutableMap.of(1, 2)).containsOnlyKeys(ImmutableSet.of(1));
+  }
+
   AbstractAssert<?, ?> testAssertThatMapContainsValue() {
     return assertThat(ImmutableMap.of(1, 2)).containsValue(3);
   }