diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java index 6045a1fc3e..b106194bac 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java @@ -15,6 +15,8 @@ */ package com.google.cloud.bigtable.data.v2.stub; +import static com.google.cloud.bigtable.data.v2.stub.CookiesHolderTest.MetadataSubject.assertThat; +import static com.google.common.truth.Truth.assertAbout; import static com.google.common.truth.Truth.assertThat; import com.google.api.gax.retrying.RetrySettings; @@ -47,6 +49,8 @@ import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.RowMutationEntry; +import com.google.common.truth.FailureMetadata; +import com.google.common.truth.Subject; import io.grpc.ForwardingServerCall; import io.grpc.Metadata; import io.grpc.MethodDescriptor; @@ -63,6 +67,7 @@ import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; +import javax.annotation.Nullable; import org.junit.After; import org.junit.Before; import org.junit.Test; @@ -149,13 +154,9 @@ public void testReadRows() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - String bytes1 = lastMetadata.get(ROUTING_COOKIE_1); - String bytes2 = lastMetadata.get(ROUTING_COOKIE_2); - assertThat(bytes1).isEqualTo("readRows"); - assertThat(bytes2).isEqualTo(testCookie); - - // make sure bad key is not added - assertThat(lastMetadata.containsKey(BAD_KEY)).isFalse(); + assertThat(lastMetadata) + .containsAtLeast(ROUTING_COOKIE_1.name(), "readRows", ROUTING_COOKIE_2.name(), testCookie); + assertThat(lastMetadata).doesNotContainKeys(BAD_KEY.name()); serverMetadata.clear(); } @@ -169,13 +170,9 @@ public void testReadRow() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - String bytes1 = lastMetadata.get(ROUTING_COOKIE_1); - String bytes2 = lastMetadata.get(ROUTING_COOKIE_2); - assertThat(bytes1).isEqualTo("readRows"); - assertThat(bytes2).isEqualTo(testCookie); - - // make sure bad key is not added - assertThat(lastMetadata.containsKey(BAD_KEY)).isFalse(); + assertThat(lastMetadata) + .containsAtLeast(ROUTING_COOKIE_1.name(), "readRows", ROUTING_COOKIE_2.name(), testCookie); + assertThat(lastMetadata).doesNotContainKeys(BAD_KEY.name()); serverMetadata.clear(); } @@ -191,13 +188,10 @@ public void testMutateRows() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - String bytes1 = lastMetadata.get(ROUTING_COOKIE_1); - String bytes2 = lastMetadata.get(ROUTING_COOKIE_2); - assertThat(bytes1).isEqualTo("mutateRows"); - assertThat(bytes2).isEqualTo(testCookie); - - // make sure bad key is not added - assertThat(lastMetadata.containsKey(BAD_KEY)).isFalse(); + assertThat(lastMetadata) + .containsAtLeast( + ROUTING_COOKIE_1.name(), "mutateRows", ROUTING_COOKIE_2.name(), testCookie); + assertThat(lastMetadata).doesNotContainKeys(BAD_KEY.name()); serverMetadata.clear(); } @@ -211,13 +205,9 @@ public void testMutateRow() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - String bytes1 = lastMetadata.get(ROUTING_COOKIE_1); - String bytes2 = lastMetadata.get(ROUTING_COOKIE_2); - assertThat(bytes1).isEqualTo("mutateRow"); - assertThat(bytes2).isEqualTo(testCookie); - - // make sure bad key is not added - assertThat(lastMetadata.containsKey(BAD_KEY)).isFalse(); + assertThat(lastMetadata) + .containsAtLeast(ROUTING_COOKIE_1.name(), "mutateRow", ROUTING_COOKIE_2.name(), testCookie); + assertThat(lastMetadata).doesNotContainKeys(BAD_KEY.name()); serverMetadata.clear(); } @@ -232,13 +222,10 @@ public void testSampleRowKeys() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - String bytes1 = lastMetadata.get(ROUTING_COOKIE_1); - String bytes2 = lastMetadata.get(ROUTING_COOKIE_2); - assertThat(bytes1).isEqualTo("sampleRowKeys"); - assertThat(bytes2).isEqualTo(testCookie); - - // make sure bad key is not added - assertThat(lastMetadata.containsKey(BAD_KEY)).isFalse(); + assertThat(lastMetadata) + .containsAtLeast( + ROUTING_COOKIE_1.name(), "sampleRowKeys", ROUTING_COOKIE_2.name(), testCookie); + assertThat(lastMetadata).doesNotContainKeys(BAD_KEY.name()); serverMetadata.clear(); } @@ -254,11 +241,8 @@ public void testNoCookieSucceedReadRows() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_1)).isFalse(); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_2)).isFalse(); + assertThat(lastMetadata).doesNotContainKeys(ROUTING_COOKIE_1.name(), ROUTING_COOKIE_2.name()); - // make sure bad key is not added - assertThat(lastMetadata.containsKey(BAD_KEY)).isFalse(); serverMetadata.clear(); } @@ -273,8 +257,8 @@ public void testNoCookieSucceedReadRow() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_1)).isFalse(); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_2)).isFalse(); + assertThat(lastMetadata) + .doesNotContainKeys(ROUTING_COOKIE_1.name(), ROUTING_COOKIE_2.name(), BAD_KEY.name()); serverMetadata.clear(); } @@ -292,8 +276,8 @@ public void testNoCookieSucceedMutateRows() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_1)).isFalse(); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_2)).isFalse(); + assertThat(lastMetadata) + .doesNotContainKeys(ROUTING_COOKIE_1.name(), ROUTING_COOKIE_2.name(), BAD_KEY.name()); serverMetadata.clear(); } @@ -309,8 +293,8 @@ public void testNoCookieSucceedMutateRow() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_1)).isFalse(); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_2)).isFalse(); + assertThat(lastMetadata) + .doesNotContainKeys(ROUTING_COOKIE_1.name(), ROUTING_COOKIE_2.name(), BAD_KEY.name()); serverMetadata.clear(); } @@ -326,8 +310,8 @@ public void testNoCookieSucceedSampleRowKeys() { Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_1)).isFalse(); - assertThat(lastMetadata.containsKey(ROUTING_COOKIE_2)).isFalse(); + assertThat(lastMetadata) + .doesNotContainKeys(ROUTING_COOKIE_1.name(), ROUTING_COOKIE_2.name(), BAD_KEY.name()); serverMetadata.clear(); } @@ -374,10 +358,10 @@ public void sendHeaders(Metadata headers) { client.readRows(Query.create("table")).iterator().hasNext(); Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1); - String bytes1 = lastMetadata.get(ROUTING_COOKIE_2); - assertThat(bytes1).isEqualTo(testCookie); - String bytes2 = lastMetadata.get(routingCookieKey); - assertThat(bytes2).isEqualTo(routingCookieValue); + + assertThat(lastMetadata) + .containsAtLeast( + ROUTING_COOKIE_2.name(), testCookie, routingCookieKey.name(), routingCookieValue); newServer.shutdown(); } @@ -568,4 +552,40 @@ private void maybePopulateCookie(Metadata trailers, String label) { } } } + + static final class MetadataSubject extends Subject { + + @Nullable private final Metadata actual; + + public static Factory metadata() { + return MetadataSubject::new; + } + + private MetadataSubject(FailureMetadata metadata, @Nullable Metadata actual) { + super(metadata, actual); + this.actual = actual; + } + + public static MetadataSubject assertThat(@Nullable Metadata actual) { + return assertAbout(metadata()).that(actual); + } + + public void containsAtLeast(String... keyValuePairs) { + assert actual != null; + for (int i = 0; i < keyValuePairs.length; i += 2) { + check("containsAtLeast()") + .that(actual.get(Metadata.Key.of(keyValuePairs[i], Metadata.ASCII_STRING_MARSHALLER))) + .isEqualTo(keyValuePairs[i + 1]); + } + } + + public void doesNotContainKeys(String... keys) { + assert actual != null; + for (String key : keys) { + check("doesNotContainKeys()") + .that(actual.containsKey(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER))) + .isFalse(); + } + } + } }