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() {