func TestVTGateStreamExecuteKeyRanges(t *testing.T) { s := createSandbox("TestVTGateStreamExecuteKeyRanges") sbc := &sandboxConn{} s.MapTestConn("-20", sbc) kr, err := key.ParseKeyRangeParts("", "20") sq := proto.KeyRangeQuery{ Sql: "query", Keyspace: "TestVTGateStreamExecuteKeyRanges", KeyRanges: []key.KeyRange{kr}, TabletType: topo.TYPE_MASTER, } // Test for successful execution var qrs []*proto.QueryResult err = RpcVTGate.StreamExecuteKeyRanges(nil, &sq, 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) proto.PopulateQueryResult(singleRowResult, row) want := []*proto.QueryResult{row} if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } sq.Session = new(proto.Session) qrs = nil RpcVTGate.Begin(nil, sq.Session) err = RpcVTGate.StreamExecuteKeyRanges(nil, &sq, 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: "TestVTGateStreamExecuteKeyRanges", Shard: "-20", TransactionId: 1, TabletType: topo.TYPE_MASTER, }}, }, }, } if !reflect.DeepEqual(want, qrs) { t.Errorf("want \n%+v, got \n%+v", want, qrs) } // Test for error condition - multiple shards kr, err = key.ParseKeyRangeParts("10", "40") sq.KeyRanges = []key.KeyRange{kr} err = RpcVTGate.StreamExecuteKeyRanges(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err == nil { t.Errorf("want not nil, got %v", err) } // Test for error condition - multiple shards, non-partial keyspace kr, err = key.ParseKeyRangeParts("", "") sq.KeyRanges = []key.KeyRange{kr} err = RpcVTGate.StreamExecuteKeyRanges(nil, &sq, func(r *proto.QueryResult) error { qrs = append(qrs, r) return nil }) if err == nil { t.Errorf("want not nil, got %v", err) } }
func TestVTGateExecuteKeyRanges(t *testing.T) { s := createSandbox("TestVTGateExecuteKeyRanges") sbc1 := &sandboxConn{} sbc2 := &sandboxConn{} s.MapTestConn("-20", sbc1) s.MapTestConn("20-40", sbc2) kr, err := key.ParseKeyRangeParts("", "20") q := proto.KeyRangeQuery{ Sql: "query", Keyspace: "TestVTGateExecuteKeyRanges", KeyRanges: []key.KeyRange{kr}, TabletType: topo.TYPE_MASTER, } // Test for successful execution qr := new(proto.QueryResult) err = RpcVTGate.ExecuteKeyRanges(nil, &q, qr) if err != nil { t.Errorf("want nil, got %v", err) } wantqr := new(proto.QueryResult) proto.PopulateQueryResult(singleRowResult, wantqr) 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) } if sbc1.ExecCount != 1 { t.Errorf("want 1, got %v\n", sbc1.ExecCount) } // Test for successful execution in transaction q.Session = new(proto.Session) RpcVTGate.Begin(nil, q.Session) if !q.Session.InTransaction { t.Errorf("want true, got false") } err = RpcVTGate.ExecuteKeyRanges(nil, &q, qr) if err != nil { t.Errorf("want nil, got %v", err) } wantSession := &proto.Session{ InTransaction: true, ShardSessions: []*proto.ShardSession{{ Keyspace: "TestVTGateExecuteKeyRanges", Shard: "-20", TransactionId: 1, TabletType: topo.TYPE_MASTER, }}, } if !reflect.DeepEqual(wantSession, q.Session) { t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session) } RpcVTGate.Commit(nil, q.Session) if sbc1.CommitCount.Get() != 1 { t.Errorf("want 1, got %v", sbc1.CommitCount.Get()) } // Test for multiple shards kr, err = key.ParseKeyRangeParts("10", "30") q.KeyRanges = []key.KeyRange{kr} RpcVTGate.ExecuteKeyRanges(nil, &q, qr) if qr.RowsAffected != 2 { t.Errorf("want 2, got %v", qr.RowsAffected) } }