Skip to content

Commit d593902

Browse files
bhshkhbhshkh
authored
test(bigtable): Correct the partial batch testcase (#11921)
* test(bigtable): Correct the partial batch testcase * test partial batch with reset
1 parent 1f06452 commit d593902

File tree

1 file changed

+101
-23
lines changed

1 file changed

+101
-23
lines changed

‎bigtable/bigtable_test.go

+101-23
Original file line numberDiff line numberDiff line change
@@ -1754,8 +1754,8 @@ func TestExecuteQuery(t *testing.T) {
17541754
newPrepareQueryResp(preparedQuery1, colFamAddress, colFamInfo),
17551755
},
17561756
mockRecvMsgResps: []recvMsgResp{
1757-
newExecQueryRespPartialBatchFirstHalf(true /* reset */, nil /* sleep */, colFamAddress),
1758-
newExecQueryRespPartialBatchSecondHalf(false /* reset */, nil /* sleep */, colFamInfo),
1757+
newExecQueryRespPartialBatchFirstHalf(true /* reset */, nil /* sleep */, []string{colFamAddress}),
1758+
newExecQueryRespPartialBatchSecondHalf(false /* reset */, nil /* sleep */, []string{colFamAddress}, []string{colFamInfo}),
17591759
newExecQueryRespResumeToken(),
17601760
{err: io.EOF},
17611761
},
@@ -1764,6 +1764,37 @@ func TestExecuteQuery(t *testing.T) {
17641764
},
17651765
wantResultRowValues: [][]*btpb.Value{newProtoRowValuesWithKey(colFamAddress, colFamInfo)},
17661766
},
1767+
{
1768+
/*
1769+
1. PrepareQuery
1770+
2. ExecuteQuery
1771+
3. RecvMsg - gets first half of a batch of data with reset true - half1
1772+
3. RecvMsg - receives Unavailable error
1773+
3. RecvMsg - gets first half of a batch of data with reset true - half2
1774+
4. RecvMsg - gets second half of a batch of data with reset false - half3
1775+
5. RecvMsg - gets resume token
1776+
6. RecvMsg - gets EOF
1777+
1778+
The resulting row should contain only half2 and half3
1779+
*/
1780+
desc: "success with single batch received across multiple responses with reset",
1781+
mockPrepQueryResps: []prepareQueryResp{
1782+
newPrepareQueryResp(preparedQuery1, colFamAddress, colFamInfo),
1783+
},
1784+
mockRecvMsgResps: []recvMsgResp{
1785+
newExecQueryRespPartialBatchFirstHalf(true /* reset */, nil /* sleep */, []string{colFamAddress}),
1786+
{err: status.Error(codes.Unavailable, "mock unavailable error")},
1787+
newExecQueryRespPartialBatchFirstHalf(true /* reset */, nil /* sleep */, []string{colFamAddressNew}),
1788+
newExecQueryRespPartialBatchSecondHalf(false /* reset */, nil /* sleep */, []string{colFamAddressNew}, []string{colFamInfo}),
1789+
newExecQueryRespResumeToken(),
1790+
{err: io.EOF},
1791+
},
1792+
wantExecReqPrepQuerys: [][]byte{
1793+
[]byte(preparedQuery1),
1794+
[]byte(preparedQuery1),
1795+
},
1796+
wantResultRowValues: [][]*btpb.Value{newProtoRowValuesWithKey(colFamAddressNew, colFamInfo)},
1797+
},
17671798
{
17681799
/*
17691800
1. PrepareQuery
@@ -1776,7 +1807,7 @@ func TestExecuteQuery(t *testing.T) {
17761807
newPrepareQueryResp(preparedQuery1, colFamAddress),
17771808
},
17781809
mockRecvMsgResps: []recvMsgResp{
1779-
newExecQueryRespFullBatch(true, nil /* sleep */, colFamAddress),
1810+
newExecQueryRespFullBatch(true, nil /* sleep */, []string{colFamAddress}),
17801811
{err: io.EOF},
17811812
},
17821813
wantExecReqPrepQuerys: [][]byte{
@@ -1802,7 +1833,7 @@ func TestExecuteQuery(t *testing.T) {
18021833
},
18031834
mockRecvMsgResps: []recvMsgResp{
18041835
{err: aePcf},
1805-
newExecQueryRespFullBatch(true, nil /* sleep */, colFamAddress),
1836+
newExecQueryRespFullBatch(true, nil /* sleep */, []string{colFamAddress}),
18061837
newExecQueryRespResumeToken(),
18071838
{err: io.EOF},
18081839
},
@@ -1827,7 +1858,7 @@ func TestExecuteQuery(t *testing.T) {
18271858
newPrepareQueryResp(preparedQuery1, colFamAddress), // PrepareStatement
18281859
},
18291860
mockRecvMsgResps: []recvMsgResp{
1830-
newExecQueryRespFullBatch(true, nil /* sleep */, colFamAddress),
1861+
newExecQueryRespFullBatch(true, nil /* sleep */, []string{colFamAddress}),
18311862
newExecQueryRespResumeToken(),
18321863
{err: status.Error(codes.Unavailable, "transient error")},
18331864
{err: io.EOF},
@@ -1855,7 +1886,7 @@ func TestExecuteQuery(t *testing.T) {
18551886
newPrepareQueryResp(preparedQuery1, colFamAddress), // From Execute, because expired query
18561887
},
18571888
mockRecvMsgResps: []recvMsgResp{
1858-
newExecQueryRespFullBatch(true, nil /* sleep */, colFamAddress),
1889+
newExecQueryRespFullBatch(true, nil /* sleep */, []string{colFamAddress}),
18591890
{
18601891
sleep: ptr(testPreparedQueryTTL + 2*time.Second),
18611892
err: status.Error(codes.DeadlineExceeded, "context deadline exceeded"), // retryable
@@ -1889,7 +1920,7 @@ func TestExecuteQuery(t *testing.T) {
18891920
},
18901921
mockRecvMsgResps: []recvMsgResp{
18911922
{err: aePcf},
1892-
newExecQueryRespFullBatch(true, nil /* sleep */, colFamAddress),
1923+
newExecQueryRespFullBatch(true, nil /* sleep */, []string{colFamAddress}),
18931924
newExecQueryRespResumeToken(),
18941925
{err: io.EOF},
18951926
},
@@ -1920,13 +1951,13 @@ func TestExecuteQuery(t *testing.T) {
19201951
newPrepareQueryResp(preparedQuery2, colFamAddress, colFamInfo), // Step 8
19211952
},
19221953
mockRecvMsgResps: []recvMsgResp{
1923-
newExecQueryRespFullBatch(true, nil /* sleep */, colFamAddress), // Step 3
1924-
newExecQueryRespFullBatch(false, nil /* sleep */, colFamAddress), // Step 4
1925-
{err: status.Error(codes.Unavailable, "mock unavailable error")}, // Step 5, retryable error
1954+
newExecQueryRespFullBatch(true, nil /* sleep */, []string{colFamAddress}), // Step 3
1955+
newExecQueryRespFullBatch(false, nil /* sleep */, []string{colFamAddress}), // Step 4
1956+
{err: status.Error(codes.Unavailable, "mock unavailable error")}, // Step 5, retryable error
19261957
{err: aePcf}, // Step 7, retryable error
1927-
newExecQueryRespFullBatch(true, nil /* sleep */, colFamAddress, colFamInfo), // Step 10
1928-
newExecQueryRespResumeToken(), // Step 11
1929-
{err: io.EOF}, // Step 12
1958+
newExecQueryRespFullBatch(true, nil /* sleep */, []string{colFamAddress, colFamInfo}), // Step 10
1959+
newExecQueryRespResumeToken(), // Step 11
1960+
{err: io.EOF}, // Step 12
19301961
},
19311962
wantExecReqPrepQuerys: [][]byte{
19321963
[]byte(preparedQuery1), // Step 2
@@ -2021,6 +2052,7 @@ func TestExecuteQuery(t *testing.T) {
20212052

20222053
const testPreparedQueryTTL = 10 * time.Second
20232054
const colFamAddress = "address"
2055+
const colFamAddressNew = "address-new" // Used only for values and not metadata
20242056
const colFamInfo = "info"
20252057

20262058
var cfToValues = map[string][]*btpb.Value{
@@ -2062,6 +2094,44 @@ var cfToValues = map[string][]*btpb.Value{
20622094
},
20632095
},
20642096
},
2097+
colFamAddressNew: {
2098+
{
2099+
Kind: &btpb.Value_ArrayValue{
2100+
ArrayValue: &btpb.ArrayValue{
2101+
Values: []*btpb.Value{
2102+
{
2103+
Kind: &btpb.Value_BytesValue{
2104+
BytesValue: []byte("city"),
2105+
},
2106+
},
2107+
{
2108+
Kind: &btpb.Value_BytesValue{
2109+
BytesValue: []byte("Kirkland"),
2110+
},
2111+
},
2112+
},
2113+
},
2114+
},
2115+
},
2116+
{
2117+
Kind: &btpb.Value_ArrayValue{
2118+
ArrayValue: &btpb.ArrayValue{
2119+
Values: []*btpb.Value{
2120+
{
2121+
Kind: &btpb.Value_BytesValue{
2122+
BytesValue: []byte("state"),
2123+
},
2124+
},
2125+
{
2126+
Kind: &btpb.Value_BytesValue{
2127+
BytesValue: []byte("WA"),
2128+
},
2129+
},
2130+
},
2131+
},
2132+
},
2133+
},
2134+
},
20652135
colFamInfo: {
20662136
{
20672137
Kind: &btpb.Value_ArrayValue{
@@ -2134,9 +2204,14 @@ func newPrepareQueryResp(preparedQuery string, colFams ...string) prepareQueryRe
21342204
}
21352205
}
21362206

2137-
func newRecvMsgResp(reset bool, protoRows *btpb.ProtoRows, blockTime *time.Duration) recvMsgResp {
2207+
func newRecvMsgResp(reset bool, protoRows, checksumProtoRows *btpb.ProtoRows, blockTime *time.Duration) recvMsgResp {
21382208
marshalled, _ := proto.Marshal(protoRows)
2139-
checksum := crc32.Checksum(marshalled, crc32cTable)
2209+
2210+
var checksum *uint32
2211+
if checksumProtoRows != nil {
2212+
marshalledCk, _ := proto.Marshal(checksumProtoRows)
2213+
checksum = ptr(crc32.Checksum(marshalledCk, crc32cTable))
2214+
}
21402215
return recvMsgResp{
21412216
results: &btpb.ExecuteQueryResponse_Results{
21422217
Results: &btpb.PartialResultSet{
@@ -2145,7 +2220,7 @@ func newRecvMsgResp(reset bool, protoRows *btpb.ProtoRows, blockTime *time.Durat
21452220
BatchData: marshalled,
21462221
},
21472222
},
2148-
BatchChecksum: &checksum,
2223+
BatchChecksum: checksum,
21492224
Reset_: reset,
21502225
},
21512226
},
@@ -2163,25 +2238,28 @@ func newExecQueryRespResumeToken() recvMsgResp {
21632238
}
21642239
}
21652240

2166-
func newExecQueryRespFullBatch(reset bool, blockTime *time.Duration, colFams ...string) recvMsgResp {
2241+
func newExecQueryRespFullBatch(reset bool, blockTime *time.Duration, colFams []string) recvMsgResp {
21672242
protoRows := &btpb.ProtoRows{
21682243
Values: newProtoRowValuesWithKey(colFams...),
21692244
}
2170-
return newRecvMsgResp(reset, protoRows, blockTime)
2245+
return newRecvMsgResp(reset, protoRows, protoRows, blockTime)
21712246
}
21722247

2173-
func newExecQueryRespPartialBatchFirstHalf(reset bool, blockTime *time.Duration, colFams ...string) recvMsgResp {
2248+
func newExecQueryRespPartialBatchFirstHalf(reset bool, blockTime *time.Duration, colFams []string) recvMsgResp {
21742249
protoRows := &btpb.ProtoRows{
21752250
Values: newProtoRowValuesWithKey(colFams...),
21762251
}
2177-
return newRecvMsgResp(reset, protoRows, blockTime)
2252+
return newRecvMsgResp(reset, protoRows, nil, blockTime)
21782253
}
21792254

2180-
func newExecQueryRespPartialBatchSecondHalf(reset bool, blockTime *time.Duration, colFams ...string) recvMsgResp {
2255+
func newExecQueryRespPartialBatchSecondHalf(reset bool, blockTime *time.Duration, colFamsFirstHalf, colFamsSecondHalf []string) recvMsgResp {
21812256
protoRows := &btpb.ProtoRows{
2182-
Values: newProtoRowValues(colFams...),
2257+
Values: newProtoRowValues(colFamsSecondHalf...),
2258+
}
2259+
checksumProtoRows := &btpb.ProtoRows{
2260+
Values: append(newProtoRowValuesWithKey(colFamsFirstHalf...), newProtoRowValues(colFamsSecondHalf...)...),
21832261
}
2184-
return newRecvMsgResp(reset, protoRows, blockTime)
2262+
return newRecvMsgResp(reset, protoRows, checksumProtoRows, blockTime)
21852263
}
21862264

21872265
func newProtoRowValuesWithKey(colFams ...string) []*btpb.Value {

0 commit comments

Comments
 (0)