From e14f8ce623370d885fd3b1b212879c3dce2c9493 Mon Sep 17 00:00:00 2001 From: Anthony Yeh Date: Tue, 22 Sep 2015 14:23:37 -0700 Subject: [PATCH] java: Move RpcClient test suite to flavor-agnostic location. --- java/client/pom.xml | 12 + .../youtube/vitess/client/RpcClientTest.java | 355 ++++++++++++++++++ java/grpc-client/pom.xml | 7 + .../vitess/client/grpc/GrpcClient.java | 10 +- .../vitess/client/grpc/GrpcClientTest.java | 340 +---------------- 5 files changed, 379 insertions(+), 345 deletions(-) create mode 100644 java/client/src/test/java/com/youtube/vitess/client/RpcClientTest.java diff --git a/java/client/pom.xml b/java/client/pom.xml index 43693e3395..af075f8bbb 100644 --- a/java/client/pom.xml +++ b/java/client/pom.xml @@ -96,6 +96,18 @@ + + org.apache.maven.plugins + maven-jar-plugin + 2.4 + + + + test-jar + + + + diff --git a/java/client/src/test/java/com/youtube/vitess/client/RpcClientTest.java b/java/client/src/test/java/com/youtube/vitess/client/RpcClientTest.java new file mode 100644 index 0000000000..1ce944ed75 --- /dev/null +++ b/java/client/src/test/java/com/youtube/vitess/client/RpcClientTest.java @@ -0,0 +1,355 @@ +package com.youtube.vitess.client; + +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; + +import com.youtube.vitess.proto.Query.QueryResult; +import com.youtube.vitess.proto.Topodata.KeyRange; +import com.youtube.vitess.proto.Topodata.KeyspaceIdType; +import com.youtube.vitess.proto.Topodata.ShardReference; +import com.youtube.vitess.proto.Topodata.SrvKeyspace; +import com.youtube.vitess.proto.Topodata.SrvKeyspace.KeyspacePartition; +import com.youtube.vitess.proto.Topodata.TabletType; +import com.youtube.vitess.proto.Vtgate.SplitQueryResponse; +import com.youtube.vitess.proto.Vtrpc.CallerID; + +import org.joda.time.Duration; +import org.junit.Assert; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * RpcClientTest tests a given implementation of RpcClient + * against a mock vtgate server (go/cmd/vtgateclienttest). + * + * Each implementation should extend this class and add a @BeforeClass method that starts the + * vtgateclienttest server with the necessary parameters, and then sets 'client'. + */ +public abstract class RpcClientTest { + protected static RpcClient client; + protected static String vtRoot; + + private Context ctx; + private VTGateConn conn; + + @BeforeClass + public static void setUpBeforeSubclass() { + vtRoot = System.getenv("VTROOT"); + if (vtRoot == null) { + throw new RuntimeException("cannot find env variable VTROOT; make sure to source dev.env"); + } + } + + @Before + public void setUp() { + ctx = Context.getDefault().withDeadlineAfter(Duration.millis(5000)).withCallerId(CALLER_ID); + conn = new VTGateConn(client); + } + + private static final String ECHO_PREFIX = "echo://"; + + private static final String QUERY = "test query"; + private static final String KEYSPACE = "test_keyspace"; + + private static final List SHARDS = Arrays.asList("-80", "80-"); + private static final String SHARDS_ECHO = "[-80 80-]"; + + private static final List KEYSPACE_IDS = + Arrays.asList(new byte[] {1, 2, 3, 4}, new byte[] {5, 6, 7, 8}); + private static final String KEYSPACE_IDS_ECHO = "[[1 2 3 4] [5 6 7 8]]"; + private static final String KEYSPACE_IDS_ECHO_OLD = "[01020304 05060708]"; + + private static final List KEY_RANGES = Arrays.asList( + KeyRange.newBuilder() + .setStart(ByteString.copyFrom(new byte[] {1, 2, 3, 4})) + .setEnd(ByteString.copyFrom(new byte[] {5, 6, 7, 8})) + .build()); + private static final String KEY_RANGES_ECHO = + "[start:\"\\001\\002\\003\\004\" end:\"\\005\\006\\007\\010\" ]"; + + private static final Map ENTITY_KEYSPACE_IDS = + new ImmutableMap.Builder() + .put(new byte[] {1, 2, 3}, 123) + .put(new byte[] {4, 5, 6}, 2.0) + .put(new byte[] {7, 8, 9}, new byte[] {1, 2, 3}) + .build(); + private static final String ENTITY_KEYSPACE_IDS_ECHO = + "[xid_type:TYPE_INT xid_int:123 keyspace_id:\"\\001\\002\\003\" xid_type:TYPE_FLOAT xid_float:2 keyspace_id:\"\\004\\005\\006\" xid_type:TYPE_BYTES xid_bytes:\"\\001\\002\\003\" keyspace_id:\"\\007\\010\\t\" ]"; + + private static final TabletType TABLET_TYPE = TabletType.REPLICA; + private static final String TABLET_TYPE_ECHO = TABLET_TYPE.toString(); + + private static final Map BIND_VARS = + new ImmutableMap.Builder() + .put("int", 123) + .put("float", 2.0) + .put("bytes", new byte[] {1, 2, 3}) + .build(); + private static final String BIND_VARS_ECHO = "map[bytes:[1 2 3] float:2 int:123]"; + + private static final String SESSION_ECHO = "InTransaction: true, ShardSession: []"; + + private static final CallerID CALLER_ID = + CallerID.newBuilder() + .setPrincipal("test_principal") + .setComponent("test_component") + .setSubcomponent("test_subcomponent") + .build(); + private static final String CALLER_ID_ECHO = + "principal:\"test_principal\" component:\"test_component\" subcomponent:\"test_subcomponent\" "; + + private static Map getEcho(QueryResult result) { + Map fields = new HashMap(); + for (int i = 0; i < result.getFieldsCount(); i++) { + fields.put(result.getFields(i).getName(), result.getRows(0).getValues(i).toStringUtf8()); + } + return fields; + } + + @Test + public void testEchoExecute() throws Exception { + Map echo; + + echo = getEcho(conn.execute(ctx, ECHO_PREFIX + QUERY, BIND_VARS, TABLET_TYPE)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho( + conn.executeShards(ctx, ECHO_PREFIX + QUERY, KEYSPACE, SHARDS, BIND_VARS, TABLET_TYPE)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho(conn.executeKeyspaceIds( + ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEYSPACE_IDS, BIND_VARS, TABLET_TYPE)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(KEYSPACE_IDS_ECHO, echo.get("keyspaceIds")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho(conn.executeKeyRanges( + ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEY_RANGES, BIND_VARS, TABLET_TYPE)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(KEY_RANGES_ECHO, echo.get("keyRanges")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho(conn.executeEntityIds(ctx, ECHO_PREFIX + QUERY, KEYSPACE, "column1", + ENTITY_KEYSPACE_IDS, BIND_VARS, TABLET_TYPE)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals("column1", echo.get("entityColumnName")); + Assert.assertEquals(ENTITY_KEYSPACE_IDS_ECHO, echo.get("entityIds")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho(conn.executeBatchShards(ctx, Arrays.asList(Proto.bindShardQuery(KEYSPACE, SHARDS, + ECHO_PREFIX + QUERY, BIND_VARS)), + TABLET_TYPE, true).get(0)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho( + conn.executeBatchKeyspaceIds(ctx, Arrays.asList(Proto.bindKeyspaceIdQuery(KEYSPACE, + KEYSPACE_IDS, ECHO_PREFIX + QUERY, BIND_VARS)), + TABLET_TYPE, true).get(0)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(KEYSPACE_IDS_ECHO_OLD, echo.get("keyspaceIds")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + } + + @Test + public void testEchoStreamExecute() throws Exception { + Map echo; + + echo = getEcho(conn.streamExecute(ctx, ECHO_PREFIX + QUERY, BIND_VARS, TABLET_TYPE).next()); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho( + conn.streamExecuteShards(ctx, ECHO_PREFIX + QUERY, KEYSPACE, SHARDS, BIND_VARS, TABLET_TYPE) + .next()); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho(conn.streamExecuteKeyspaceIds(ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEYSPACE_IDS, + BIND_VARS, TABLET_TYPE).next()); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(KEYSPACE_IDS_ECHO, echo.get("keyspaceIds")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + + echo = getEcho(conn.streamExecuteKeyRanges(ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEY_RANGES, + BIND_VARS, TABLET_TYPE).next()); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(KEY_RANGES_ECHO, echo.get("keyRanges")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + } + + @Test + public void testEchoTransactionExecute() throws Exception { + Map echo; + + VTGateTx tx = conn.begin(ctx); + + echo = getEcho(tx.execute(ctx, ECHO_PREFIX + QUERY, BIND_VARS, TABLET_TYPE, true)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + Assert.assertEquals(SESSION_ECHO, echo.get("session")); + Assert.assertEquals("true", echo.get("notInTransaction")); + + echo = getEcho( + tx.executeShards(ctx, ECHO_PREFIX + QUERY, KEYSPACE, SHARDS, BIND_VARS, TABLET_TYPE, true)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + Assert.assertEquals(SESSION_ECHO, echo.get("session")); + Assert.assertEquals("true", echo.get("notInTransaction")); + + echo = getEcho(tx.executeKeyspaceIds( + ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEYSPACE_IDS, BIND_VARS, TABLET_TYPE, true)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(KEYSPACE_IDS_ECHO, echo.get("keyspaceIds")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + Assert.assertEquals(SESSION_ECHO, echo.get("session")); + Assert.assertEquals("true", echo.get("notInTransaction")); + + echo = getEcho(tx.executeKeyRanges( + ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEY_RANGES, BIND_VARS, TABLET_TYPE, true)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(KEY_RANGES_ECHO, echo.get("keyRanges")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + Assert.assertEquals(SESSION_ECHO, echo.get("session")); + Assert.assertEquals("true", echo.get("notInTransaction")); + + echo = getEcho(tx.executeEntityIds(ctx, ECHO_PREFIX + QUERY, KEYSPACE, "column1", + ENTITY_KEYSPACE_IDS, BIND_VARS, TABLET_TYPE, true)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals("column1", echo.get("entityColumnName")); + Assert.assertEquals(ENTITY_KEYSPACE_IDS_ECHO, echo.get("entityIds")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + Assert.assertEquals(SESSION_ECHO, echo.get("session")); + Assert.assertEquals("true", echo.get("notInTransaction")); + + tx.rollback(ctx); + tx = conn.begin(ctx); + + echo = getEcho(tx.executeBatchShards(ctx, Arrays.asList(Proto.bindShardQuery(KEYSPACE, SHARDS, + ECHO_PREFIX + QUERY, BIND_VARS)), + TABLET_TYPE, true).get(0)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + Assert.assertEquals(SESSION_ECHO, echo.get("session")); + + echo = + getEcho(tx.executeBatchKeyspaceIds(ctx, Arrays.asList(Proto.bindKeyspaceIdQuery(KEYSPACE, + KEYSPACE_IDS, ECHO_PREFIX + QUERY, BIND_VARS)), + TABLET_TYPE, true).get(0)); + Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); + Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); + Assert.assertEquals(KEYSPACE, echo.get("keyspace")); + Assert.assertEquals(KEYSPACE_IDS_ECHO_OLD, echo.get("keyspaceIds")); + Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); + Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); + Assert.assertEquals(SESSION_ECHO, echo.get("session")); + + tx.commit(ctx); + } + + @Test + public void testEchoSplitQuery() throws Exception { + SplitQueryResponse.Part expected = + SplitQueryResponse.Part.newBuilder() + .setQuery(Proto.bindQuery(ECHO_PREFIX + QUERY + ":split_column:123", BIND_VARS)) + .setKeyRangePart( + SplitQueryResponse.KeyRangePart.newBuilder().setKeyspace(KEYSPACE).build()) + .build(); + SplitQueryResponse.Part actual = + conn.splitQuery(ctx, KEYSPACE, ECHO_PREFIX + QUERY, BIND_VARS, "split_column", 123).get(0); + Assert.assertEquals(expected, actual); + } + + @Test + public void testGetSrvKeyspace() throws Exception { + SrvKeyspace expected = + SrvKeyspace.newBuilder() + .addPartitions( + KeyspacePartition.newBuilder() + .setServedType(TabletType.REPLICA) + .addShardReferences( + ShardReference.newBuilder() + .setName("shard0") + .setKeyRange( + KeyRange.newBuilder() + .setStart( + ByteString.copyFrom(new byte[] {0x40, 0, 0, 0, 0, 0, 0, 0})) + .setEnd(ByteString.copyFrom( + new byte[] {(byte) 0x80, 0, 0, 0, 0, 0, 0, 0})) + .build()) + .build()) + .build()) + .setShardingColumnName("sharding_column_name") + .setShardingColumnType(KeyspaceIdType.UINT64) + .addServedFrom( + SrvKeyspace.ServedFrom.newBuilder() + .setTabletType(TabletType.MASTER) + .setKeyspace("other_keyspace") + .build()) + .setSplitShardCount(128) + .build(); + SrvKeyspace actual = conn.getSrvKeyspace(ctx, "big"); + Assert.assertEquals(expected, actual); + } +} diff --git a/java/grpc-client/pom.xml b/java/grpc-client/pom.xml index ace1435aaa..935710d4fe 100644 --- a/java/grpc-client/pom.xml +++ b/java/grpc-client/pom.xml @@ -16,6 +16,13 @@ client 1.0-SNAPSHOT + + com.youtube.vitess + client + 1.0-SNAPSHOT + test-jar + test + io.grpc grpc-all diff --git a/java/grpc-client/src/main/java/com/youtube/vitess/client/grpc/GrpcClient.java b/java/grpc-client/src/main/java/com/youtube/vitess/client/grpc/GrpcClient.java index f98d6df25d..ff243836d5 100644 --- a/java/grpc-client/src/main/java/com/youtube/vitess/client/grpc/GrpcClient.java +++ b/java/grpc-client/src/main/java/com/youtube/vitess/client/grpc/GrpcClient.java @@ -1,7 +1,6 @@ package com.youtube.vitess.client.grpc; import com.youtube.vitess.client.Context; -import com.youtube.vitess.client.Proto; import com.youtube.vitess.client.RpcClient; import com.youtube.vitess.client.StreamIterator; import com.youtube.vitess.client.VitessException; @@ -218,8 +217,7 @@ public class GrpcClient implements RpcClient { } @Override - public BeginResponse begin(Context ctx, BeginRequest request) - throws VitessRpcException { + public BeginResponse begin(Context ctx, BeginRequest request) throws VitessRpcException { try (GrpcContext gctx = new GrpcContext(ctx)) { return blockingStub.begin(request); } catch (Exception e) { @@ -228,8 +226,7 @@ public class GrpcClient implements RpcClient { } @Override - public CommitResponse commit(Context ctx, CommitRequest request) - throws VitessRpcException { + public CommitResponse commit(Context ctx, CommitRequest request) throws VitessRpcException { try (GrpcContext gctx = new GrpcContext(ctx)) { return blockingStub.commit(request); } catch (Exception e) { @@ -238,8 +235,7 @@ public class GrpcClient implements RpcClient { } @Override - public RollbackResponse rollback(Context ctx, RollbackRequest request) - throws VitessRpcException { + public RollbackResponse rollback(Context ctx, RollbackRequest request) throws VitessRpcException { try (GrpcContext gctx = new GrpcContext(ctx)) { return blockingStub.rollback(request); } catch (Exception e) { diff --git a/java/grpc-client/src/test/java/com/youtube/vitess/client/grpc/GrpcClientTest.java b/java/grpc-client/src/test/java/com/youtube/vitess/client/grpc/GrpcClientTest.java index 90bef59373..b24e810837 100644 --- a/java/grpc-client/src/test/java/com/youtube/vitess/client/grpc/GrpcClientTest.java +++ b/java/grpc-client/src/test/java/com/youtube/vitess/client/grpc/GrpcClientTest.java @@ -1,52 +1,25 @@ package com.youtube.vitess.client.grpc; -import com.google.common.collect.ImmutableMap; -import com.google.protobuf.ByteString; - import com.youtube.vitess.client.Context; -import com.youtube.vitess.client.Proto; -import com.youtube.vitess.client.RpcClient; -import com.youtube.vitess.client.VTGateConn; -import com.youtube.vitess.client.VTGateTx; -import com.youtube.vitess.proto.Query.QueryResult; -import com.youtube.vitess.proto.Topodata.KeyRange; -import com.youtube.vitess.proto.Topodata.KeyspaceIdType; -import com.youtube.vitess.proto.Topodata.ShardReference; -import com.youtube.vitess.proto.Topodata.SrvKeyspace; -import com.youtube.vitess.proto.Topodata.SrvKeyspace.KeyspacePartition; -import com.youtube.vitess.proto.Topodata.TabletType; -import com.youtube.vitess.proto.Vtgate.SplitQueryResponse; -import com.youtube.vitess.proto.Vtrpc.CallerID; +import com.youtube.vitess.client.RpcClientTest; import org.joda.time.Duration; import org.junit.AfterClass; -import org.junit.Assert; -import org.junit.Before; import org.junit.BeforeClass; -import org.junit.Test; import java.net.InetSocketAddress; import java.net.ServerSocket; import java.util.Arrays; -import java.util.HashMap; -import java.util.List; -import java.util.Map; /** * This tests GrpcClient with a mock vtgate server (go/cmd/vtgateclienttest). */ -public class GrpcClientTest { +public class GrpcClientTest extends RpcClientTest { private static Process vtgateclienttest; private static int port; - private static RpcClient client; @BeforeClass public static void setUpBeforeClass() throws Exception { - String vtRoot = System.getenv("VTROOT"); - if (vtRoot == null) { - throw new RuntimeException("cannot find env variable VTROOT; make sure to source dev.env"); - } - ServerSocket socket = new ServerSocket(0); port = socket.getLocalPort(); socket.close(); @@ -70,313 +43,4 @@ public class GrpcClientTest { vtgateclienttest.destroy(); } } - - private Context ctx; - private VTGateConn conn; - - @Before - public void setUp() { - ctx = Context.getDefault().withDeadlineAfter(Duration.millis(5000)).withCallerId(CALLER_ID); - conn = new VTGateConn(client); - } - - private static final String ECHO_PREFIX = "echo://"; - - private static final String QUERY = "test query"; - private static final String KEYSPACE = "test_keyspace"; - - private static final List SHARDS = Arrays.asList("-80", "80-"); - private static final String SHARDS_ECHO = "[-80 80-]"; - - private static final List KEYSPACE_IDS = - Arrays.asList(new byte[] {1, 2, 3, 4}, new byte[] {5, 6, 7, 8}); - private static final String KEYSPACE_IDS_ECHO = "[[1 2 3 4] [5 6 7 8]]"; - private static final String KEYSPACE_IDS_ECHO_OLD = "[01020304 05060708]"; - - private static final List KEY_RANGES = Arrays.asList( - KeyRange.newBuilder() - .setStart(ByteString.copyFrom(new byte[] {1, 2, 3, 4})) - .setEnd(ByteString.copyFrom(new byte[] {5, 6, 7, 8})) - .build()); - private static final String KEY_RANGES_ECHO = "[start:\"\\001\\002\\003\\004\" end:\"\\005\\006\\007\\010\" ]"; - - private static final Map ENTITY_KEYSPACE_IDS = - new ImmutableMap.Builder() - .put(new byte[] {1, 2, 3}, 123) - .put(new byte[] {4, 5, 6}, 2.0) - .put(new byte[] {7, 8, 9}, new byte[] {1, 2, 3}) - .build(); - private static final String ENTITY_KEYSPACE_IDS_ECHO = - "[xid_type:TYPE_INT xid_int:123 keyspace_id:\"\\001\\002\\003\" xid_type:TYPE_FLOAT xid_float:2 keyspace_id:\"\\004\\005\\006\" xid_type:TYPE_BYTES xid_bytes:\"\\001\\002\\003\" keyspace_id:\"\\007\\010\\t\" ]"; - - private static final TabletType TABLET_TYPE = TabletType.REPLICA; - private static final String TABLET_TYPE_ECHO = TABLET_TYPE.toString(); - - private static final Map BIND_VARS = - new ImmutableMap.Builder() - .put("int", 123) - .put("float", 2.0) - .put("bytes", new byte[] {1, 2, 3}) - .build(); - private static final String BIND_VARS_ECHO = "map[bytes:[1 2 3] float:2 int:123]"; - - private static final String SESSION_ECHO = "InTransaction: true, ShardSession: []"; - - private static final CallerID CALLER_ID = - CallerID.newBuilder() - .setPrincipal("test_principal") - .setComponent("test_component") - .setSubcomponent("test_subcomponent") - .build(); - private static final String CALLER_ID_ECHO = - "principal:\"test_principal\" component:\"test_component\" subcomponent:\"test_subcomponent\" "; - - private static Map getEcho(QueryResult result) { - Map fields = new HashMap(); - for (int i = 0; i < result.getFieldsCount(); i++) { - fields.put(result.getFields(i).getName(), result.getRows(0).getValues(i).toStringUtf8()); - } - return fields; - } - - @Test - public void testEchoExecute() throws Exception { - Map echo; - - echo = getEcho(conn.execute(ctx, ECHO_PREFIX + QUERY, BIND_VARS, TABLET_TYPE)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho( - conn.executeShards(ctx, ECHO_PREFIX + QUERY, KEYSPACE, SHARDS, BIND_VARS, TABLET_TYPE)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho(conn.executeKeyspaceIds( - ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEYSPACE_IDS, BIND_VARS, TABLET_TYPE)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(KEYSPACE_IDS_ECHO, echo.get("keyspaceIds")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho(conn.executeKeyRanges( - ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEY_RANGES, BIND_VARS, TABLET_TYPE)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(KEY_RANGES_ECHO, echo.get("keyRanges")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho(conn.executeEntityIds(ctx, ECHO_PREFIX + QUERY, KEYSPACE, "column1", - ENTITY_KEYSPACE_IDS, BIND_VARS, TABLET_TYPE)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals("column1", echo.get("entityColumnName")); - Assert.assertEquals(ENTITY_KEYSPACE_IDS_ECHO, echo.get("entityIds")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho(conn.executeBatchShards(ctx, Arrays.asList(Proto.bindShardQuery(KEYSPACE, SHARDS, - ECHO_PREFIX + QUERY, BIND_VARS)), - TABLET_TYPE, true).get(0)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho( - conn.executeBatchKeyspaceIds(ctx, Arrays.asList(Proto.bindKeyspaceIdQuery(KEYSPACE, - KEYSPACE_IDS, ECHO_PREFIX + QUERY, BIND_VARS)), - TABLET_TYPE, true).get(0)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(KEYSPACE_IDS_ECHO_OLD, echo.get("keyspaceIds")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - } - - @Test - public void testEchoStreamExecute() throws Exception { - Map echo; - - echo = getEcho(conn.streamExecute(ctx, ECHO_PREFIX + QUERY, BIND_VARS, TABLET_TYPE).next()); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho( - conn.streamExecuteShards(ctx, ECHO_PREFIX + QUERY, KEYSPACE, SHARDS, BIND_VARS, TABLET_TYPE) - .next()); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho(conn.streamExecuteKeyspaceIds(ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEYSPACE_IDS, - BIND_VARS, TABLET_TYPE).next()); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(KEYSPACE_IDS_ECHO, echo.get("keyspaceIds")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - - echo = getEcho(conn.streamExecuteKeyRanges(ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEY_RANGES, - BIND_VARS, TABLET_TYPE).next()); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(KEY_RANGES_ECHO, echo.get("keyRanges")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - } - - @Test - public void testEchoTransactionExecute() throws Exception { - Map echo; - - VTGateTx tx = conn.begin(ctx); - - echo = getEcho(tx.execute(ctx, ECHO_PREFIX + QUERY, BIND_VARS, TABLET_TYPE, true)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - Assert.assertEquals(SESSION_ECHO, echo.get("session")); - Assert.assertEquals("true", echo.get("notInTransaction")); - - echo = getEcho( - tx.executeShards(ctx, ECHO_PREFIX + QUERY, KEYSPACE, SHARDS, BIND_VARS, TABLET_TYPE, true)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - Assert.assertEquals(SESSION_ECHO, echo.get("session")); - Assert.assertEquals("true", echo.get("notInTransaction")); - - echo = getEcho(tx.executeKeyspaceIds( - ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEYSPACE_IDS, BIND_VARS, TABLET_TYPE, true)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(KEYSPACE_IDS_ECHO, echo.get("keyspaceIds")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - Assert.assertEquals(SESSION_ECHO, echo.get("session")); - Assert.assertEquals("true", echo.get("notInTransaction")); - - echo = getEcho(tx.executeKeyRanges( - ctx, ECHO_PREFIX + QUERY, KEYSPACE, KEY_RANGES, BIND_VARS, TABLET_TYPE, true)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(KEY_RANGES_ECHO, echo.get("keyRanges")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - Assert.assertEquals(SESSION_ECHO, echo.get("session")); - Assert.assertEquals("true", echo.get("notInTransaction")); - - echo = getEcho(tx.executeEntityIds(ctx, ECHO_PREFIX + QUERY, KEYSPACE, "column1", - ENTITY_KEYSPACE_IDS, BIND_VARS, TABLET_TYPE, true)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals("column1", echo.get("entityColumnName")); - Assert.assertEquals(ENTITY_KEYSPACE_IDS_ECHO, echo.get("entityIds")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - Assert.assertEquals(SESSION_ECHO, echo.get("session")); - Assert.assertEquals("true", echo.get("notInTransaction")); - - tx.rollback(ctx); - tx = conn.begin(ctx); - - echo = getEcho(tx.executeBatchShards(ctx, Arrays.asList(Proto.bindShardQuery(KEYSPACE, SHARDS, - ECHO_PREFIX + QUERY, BIND_VARS)), - TABLET_TYPE, true).get(0)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(SHARDS_ECHO, echo.get("shards")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - Assert.assertEquals(SESSION_ECHO, echo.get("session")); - - echo = - getEcho(tx.executeBatchKeyspaceIds(ctx, Arrays.asList(Proto.bindKeyspaceIdQuery(KEYSPACE, - KEYSPACE_IDS, ECHO_PREFIX + QUERY, BIND_VARS)), - TABLET_TYPE, true).get(0)); - Assert.assertEquals(CALLER_ID_ECHO, echo.get("callerId")); - Assert.assertEquals(ECHO_PREFIX + QUERY, echo.get("query")); - Assert.assertEquals(KEYSPACE, echo.get("keyspace")); - Assert.assertEquals(KEYSPACE_IDS_ECHO_OLD, echo.get("keyspaceIds")); - Assert.assertEquals(BIND_VARS_ECHO, echo.get("bindVars")); - Assert.assertEquals(TABLET_TYPE_ECHO, echo.get("tabletType")); - Assert.assertEquals(SESSION_ECHO, echo.get("session")); - - tx.commit(ctx); - } - - @Test - public void testEchoSplitQuery() throws Exception { - SplitQueryResponse.Part expected = - SplitQueryResponse.Part.newBuilder() - .setQuery(Proto.bindQuery(ECHO_PREFIX + QUERY + ":split_column:123", BIND_VARS)) - .setKeyRangePart( - SplitQueryResponse.KeyRangePart.newBuilder().setKeyspace(KEYSPACE).build()) - .build(); - SplitQueryResponse.Part actual = - conn.splitQuery(ctx, KEYSPACE, ECHO_PREFIX + QUERY, BIND_VARS, "split_column", 123).get(0); - Assert.assertEquals(expected, actual); - } - - @Test - public void testGetSrvKeyspace() throws Exception { - SrvKeyspace expected = - SrvKeyspace.newBuilder() - .addPartitions( - KeyspacePartition.newBuilder() - .setServedType(TabletType.REPLICA) - .addShardReferences( - ShardReference.newBuilder() - .setName("shard0") - .setKeyRange( - KeyRange.newBuilder() - .setStart( - ByteString.copyFrom(new byte[] {0x40, 0, 0, 0, 0, 0, 0, 0})) - .setEnd(ByteString.copyFrom( - new byte[] {(byte) 0x80, 0, 0, 0, 0, 0, 0, 0})) - .build()) - .build()) - .build()) - .setShardingColumnName("sharding_column_name") - .setShardingColumnType(KeyspaceIdType.UINT64) - .addServedFrom( - SrvKeyspace.ServedFrom.newBuilder() - .setTabletType(TabletType.MASTER) - .setKeyspace("other_keyspace") - .build()) - .setSplitShardCount(128) - .build(); - SrvKeyspace actual = conn.getSrvKeyspace(ctx, "big"); - Assert.assertEquals(expected, actual); - } }