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..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,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 after(Map map, K key) { } } + static final class AssertThatMapContainsOnlyKeys { + @BeforeTemplate + AbstractCollectionAssert, K, ?> before(Map map, Set keys) { + return assertThat(map.keySet()).hasSameElementsAs(keys); + } + + @AfterTemplate + @UseImportPolicy(STATIC_IMPORT_ALWAYS) + MapAssert after(Map map, Set keys) { + return assertThat(map).containsOnlyKeys(keys); + } + } + static final class AssertThatMapContainsValue { @BeforeTemplate AbstractBooleanAssert before(Map 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..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 @@ -130,6 +130,10 @@ MapAssert 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/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..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 @@ -130,6 +130,10 @@ MapAssert 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); }