From 1fbff8cbbe38b0213b4054d4ddc44b999c2e52bf Mon Sep 17 00:00:00 2001
From: mohamedsamehsalah <mohamedsamehsalah.mohamed@teampicnic.com>
Date: Mon, 10 Apr 2023 12:33:11 +0200
Subject: [PATCH 1/2] Prefer usages of `containsOnlyKeys` when asserting
 keySets of a map

---
 .../errorprone/refasterrules/AssertJMapRules.java  | 14 ++++++++++++++
 .../refasterrules/AssertJMapRulesTestInput.java    |  4 ++++
 .../refasterrules/AssertJMapRulesTestOutput.java   |  4 ++++
 3 files changed, 22 insertions(+)

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
index 4f3d782a7f..a95bf0c32e 100644
--- 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/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java
index 82436281ab..b52b3a7939 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java
+++ b/error-prone-contrib/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/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java
index e6087c8267..5ad04f2376 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java
+++ b/error-prone-contrib/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);
   }

From 23fe7c651204dbda045f33d8ff177fe66af4e216 Mon Sep 17 00:00:00 2001
From: Stephan Schroevers <stephan.schroevers@teampicnic.com>
Date: Mon, 10 Apr 2023 19:20:03 +0200
Subject: [PATCH 2/2] Suggestions

---
 .../picnic/errorprone/refasterrules/AssertJMapRules.java    | 6 ++++--
 .../errorprone/refasterrules/AssertJMapRulesTestInput.java  | 2 +-
 .../errorprone/refasterrules/AssertJMapRulesTestOutput.java | 2 +-
 3 files changed, 6 insertions(+), 4 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
index a95bf0c32e..f206791760 100644
--- 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,6 +16,7 @@
 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;
@@ -26,6 +27,7 @@
 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;
@@ -223,13 +225,13 @@ 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) {
+    AbstractCollectionAssert<?, Collection<? extends K>, K, ?> 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) {
+    MapAssert<K, V> after(Map<K, V> map, Set<K> keys) {
       return assertThat(map).containsOnlyKeys(keys);
     }
   }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java
index b52b3a7939..66ebdc0ef1 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestInput.java
@@ -131,7 +131,7 @@ MapAssert<Integer, Integer> testAbstractMapAssertHasSameSizeAs() {
   }
 
   AbstractAssert<?, ?> testAssertThatMapContainsOnlyKeys() {
-    return assertThat(ImmutableMap.of(1, 2).keySet()).hasSameElementsAs(ImmutableSet.of(1));
+    return assertThat(ImmutableMap.of(1, 2).keySet()).hasSameElementsAs(ImmutableSet.of(3));
   }
 
   AbstractAssert<?, ?> testAssertThatMapContainsValue() {
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java
index 5ad04f2376..413d8fbec1 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/AssertJMapRulesTestOutput.java
@@ -131,7 +131,7 @@ MapAssert<Integer, Integer> testAbstractMapAssertHasSameSizeAs() {
   }
 
   AbstractAssert<?, ?> testAssertThatMapContainsOnlyKeys() {
-    return assertThat(ImmutableMap.of(1, 2)).containsOnlyKeys(ImmutableSet.of(1));
+    return assertThat(ImmutableMap.of(1, 2)).containsOnlyKeys(ImmutableSet.of(3));
   }
 
   AbstractAssert<?, ?> testAssertThatMapContainsValue() {