Skip to content

Commit

Permalink
clean up test
Browse files Browse the repository at this point in the history
  • Loading branch information
@mutianf
mutianf committed Nov 22, 2023
1 parent d49ce70 commit 6479464
Showing 1 changed file with 71 additions and 51 deletions.
122 changes: 71 additions & 51 deletions ...loud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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();
}
Expand All @@ -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();
}
Expand All @@ -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();
}
Expand All @@ -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();
}
Expand All @@ -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();
}
Expand All @@ -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();
}

Expand All @@ -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();
}
Expand All @@ -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();
}
Expand All @@ -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();
}
Expand All @@ -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();
}
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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<MetadataSubject, Metadata> 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();
}
}
}
}

0 comments on commit 6479464

Please sign in to comment.