Introduce AssertThatMapContainsOnlyKeys Refaster rule (#576)

PicnicSupermarket · Apr 18, 2023 · ebd64c1 · ebd64c1
1 parent 929f1dd
commit ebd64c1
Show file tree

Hide file tree

Showing 3 changed files with 24 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
@@ -16,15 +16,18 @@
 import com.google.errorprone.refaster.annotation.BeforeTemplate;
 import com.google.errorprone.refaster.annotation.UseImportPolicy;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
 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;
 import org.assertj.core.api.AbstractBooleanAssert;
+import org.assertj.core.api.AbstractCollectionAssert;
 import org.assertj.core.api.AbstractMapAssert;
 import org.assertj.core.api.MapAssert;
 import tech.picnic.errorprone.refaster.annotation.OnlineDocumentation;
@@ -220,6 +223,19 @@ MapAssert<K, V> after(Map<K, V> map, K key) {
     }
   }
 
+  static final class AssertThatMapContainsOnlyKeys<K, V> {
+    @BeforeTemplate
+    AbstractCollectionAssert<?, Collection<? extends K>, K, ?> before(Map<K, V> map, Set<K> keys) {
+      return assertThat(map.keySet()).hasSameElementsAs(keys);
+    }
+
+    @AfterTemplate
+    @UseImportPolicy(STATIC_IMPORT_ALWAYS)
+    MapAssert<K, V> 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(3));
+  }
+
   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(3));
+  }
+
   AbstractAssert<?, ?> testAssertThatMapContainsValue() {
     return assertThat(ImmutableMap.of(1, 2)).containsValue(3);
   }