func TestVTGateExecuteShard(t *testing.T) {
sandbox := createSandbox("TestVTGateExecuteShard")
sbc := &sandboxConn{}
sandbox.MapTestConn("0", sbc)
q := proto.QueryShard{
Sql: "query",
Keyspace: "TestVTGateExecuteShard",
Shards: []string{"0"},
}
qr := new(proto.QueryResult)
err := RpcVTGate.ExecuteShard(&context.DummyContext{}, &q, qr)
if err != nil {
t.Errorf("want nil, got %v", err)
}
wantqr := new(proto.QueryResult)
wantqr.Result = singleRowResult
if !reflect.DeepEqual(wantqr, qr) {
t.Errorf("want \n%+v, got \n%+v", singleRowResult, qr)
}
if qr.Session != nil {
t.Errorf("want nil, got %+v\n", qr.Session)
}
q.Session = new(proto.Session)
RpcVTGate.Begin(&context.DummyContext{}, q.Session)
if !q.Session.InTransaction {
t.Errorf("want true, got false")
}
RpcVTGate.ExecuteShard(&context.DummyContext{}, &q, qr)
wantSession := &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: "TestVTGateExecuteShard",
Shard: "0",
TransactionId: 1,
}},
}
if !reflect.DeepEqual(wantSession, q.Session) {
t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session)
}
RpcVTGate.Commit(&context.DummyContext{}, q.Session)
if sbc.CommitCount != 1 {
t.Errorf("want 1, got %d", sbc.CommitCount)
}
q.Session = new(proto.Session)
RpcVTGate.Begin(&context.DummyContext{}, q.Session)
RpcVTGate.ExecuteShard(&context.DummyContext{}, &q, qr)
RpcVTGate.Rollback(&context.DummyContext{}, q.Session)
/*
// Flaky: This test should be run manually.
runtime.Gosched()
if sbc.RollbackCount != 1 {
t.Errorf("want 1, got %d", sbc.RollbackCount)
}
*/
}
func TestVTGateStreamExecuteShard(t *testing.T) {
s := createSandbox("TestVTGateStreamExecuteShard")
sbc := &sandboxConn{}
s.MapTestConn("0", sbc)
q := proto.QueryShard{
Sql: "query",
Keyspace: "TestVTGateStreamExecuteShard",
Shards: []string{"0"},
TabletType: topo.TYPE_MASTER,
}
// Test for successful execution
var qrs []*proto.QueryResult
err := RpcVTGate.StreamExecuteShard(&context.DummyContext{}, &q, func(r *proto.QueryResult) error {
qrs = append(qrs, r)
return nil
})
if err != nil {
t.Errorf("want nil, got %v", err)
}
row := new(proto.QueryResult)
row.Result = singleRowResult
want := []*proto.QueryResult{row}
if !reflect.DeepEqual(want, qrs) {
t.Errorf("want \n%+v, got \n%+v", want, qrs)
}
q.Session = new(proto.Session)
qrs = nil
RpcVTGate.Begin(&context.DummyContext{}, q.Session)
err = RpcVTGate.StreamExecuteShard(&context.DummyContext{}, &q, func(r *proto.QueryResult) error {
qrs = append(qrs, r)
return nil
})
want = []*proto.QueryResult{
row,
&proto.QueryResult{
Session: &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: "TestVTGateStreamExecuteShard",
Shard: "0",
TransactionId: 1,
TabletType: topo.TYPE_MASTER,
}},
},
},
}
if !reflect.DeepEqual(want, qrs) {
t.Errorf("want \n%+v, got \n%+v", want, qrs)
}
}