add test helper to check the right count

Signed-off-by: Michael Demmer <mdemmer@slack-corp.com>

go/vt/vtgate/tabletgateway_test.go

@@ -34,14 +34,18 @@ import (
  vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
  "vitess.io/vitess/go/vt/topo"
  "vitess.io/vitess/go/vt/vterrors"
+ "vitess.io/vitess/go/vt/vttablet/sandboxconn"
 )
 
 func TestTabletGatewayExecute(t *testing.T) {
  ctx := utils.LeakCheckContext(t)
  testTabletGatewayGeneric(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
  _, err := tg.Execute(ctx, target, "query", nil, 0, 0, nil)
  return err
- })
+ },
+ func(t *testing.T, sc *sandboxconn.SandboxConn, want int64) {
+ assert.Equal(t, want, sc.ExecCount.Load())
+ })
  testTabletGatewayTransact(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
  _, err := tg.Execute(ctx, target, "query", nil, 1, 0, nil)
  return err
@@ -55,15 +59,21 @@ func TestTabletGatewayExecuteStream(t *testing.T) {
  return nil
  })
  return err
- })
+ },
+ func(t *testing.T, sc *sandboxconn.SandboxConn, want int64) {
+ assert.Equal(t, want, sc.ExecCount.Load())
+ })
 }
 
 func TestTabletGatewayBegin(t *testing.T) {
  ctx := utils.LeakCheckContext(t)
  testTabletGatewayGeneric(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
  _, err := tg.Begin(ctx, target, nil)
  return err
- })
+ },
+ func(t *testing.T, sc *sandboxconn.SandboxConn, want int64) {
+ assert.Equal(t, want, sc.BeginCount.Load())
+ })
 }
 
 func TestTabletGatewayCommit(t *testing.T) {
@@ -87,7 +97,11 @@ func TestTabletGatewayBeginExecute(t *testing.T) {
  testTabletGatewayGeneric(t, ctx, func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error {
  _, _, err := tg.BeginExecute(ctx, target, nil, "query", nil, 0, nil)
  return err
- })
+ },
+ func(t *testing.T, sc *sandboxconn.SandboxConn, want int64) {
+ t.Helper()
+ assert.Equal(t, want, sc.BeginCount.Load())
+ })
 }
 
 func TestTabletGatewayShuffleTablets(t *testing.T) {
@@ -177,20 +191,20 @@ func TestTabletGatewayReplicaTransactionError(t *testing.T) {
  verifyContainsError(t, err, "query service can only be used for non-transactional queries on replicas", vtrpcpb.Code_INTERNAL)
 }
 
-func testTabletGatewayGeneric(t *testing.T, ctx context.Context, f func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error) {
-
- testTabletGatewayGenericHelper(t, ctx, f)
+func testTabletGatewayGeneric(t *testing.T, ctx context.Context, f func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error, verifyExpectedCount func(t *testing.T, sc *sandboxconn.SandboxConn, want int64)) {
+ t.Helper()
+ testTabletGatewayGenericHelper(t, ctx, f, verifyExpectedCount)
 
  // test again with the balancer enabled assuming vtgates in both cells where there
  // are tablets, so that it will still route to the local cell always, but this way
  // it will test both implementations of skipping invalid tablets for retry
  balancerEnabled = true
  balancerVtgateCells = []string{"cell", "cell2"}
- testTabletGatewayGenericHelper(t, ctx, f)
+ testTabletGatewayGenericHelper(t, ctx, f, verifyExpectedCount)
  balancerEnabled = false
 }
 
-func testTabletGatewayGenericHelper(t *testing.T, ctx context.Context, f func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error) {
+func testTabletGatewayGenericHelper(t *testing.T, ctx context.Context, f func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error, verifyExpectedCount func(t *testing.T, sc *sandboxconn.SandboxConn, want int64)) {
  t.Helper()
  keyspace := "ks"
  shard := "0"
@@ -231,8 +245,8 @@ func testTabletGatewayGenericHelper(t *testing.T, ctx context.Context, f func(ct
  sc2.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
  err = f(ctx, tg, target)
  verifyContainsError(t, err, "target: ks.0.replica", vtrpcpb.Code_FAILED_PRECONDITION)
- assert.Equal(t, int64(1), sc1.ExecCount.Load())
- assert.Equal(t, int64(1), sc2.ExecCount.Load())
+ verifyExpectedCount(t, sc1, 1)
+ verifyExpectedCount(t, sc2, 1)
 
  // fatal error
  hc.Reset()
@@ -242,8 +256,8 @@ func testTabletGatewayGenericHelper(t *testing.T, ctx context.Context, f func(ct
  sc2.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
  err = f(ctx, tg, target)
  verifyContainsError(t, err, "target: ks.0.replica", vtrpcpb.Code_FAILED_PRECONDITION)
- assert.Equal(t, int64(1), sc1.ExecCount.Load())
- assert.Equal(t, int64(1), sc2.ExecCount.Load())
+ verifyExpectedCount(t, sc1, 1)
+ verifyExpectedCount(t, sc2, 1)
 
  // server error - no retry
  hc.Reset()
@@ -252,17 +266,17 @@ func testTabletGatewayGenericHelper(t *testing.T, ctx context.Context, f func(ct
  sc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
  err = f(ctx, tg, target)
  assert.Equal(t, vtrpcpb.Code_INVALID_ARGUMENT, vterrors.Code(err))
- assert.Equal(t, int64(1), sc1.ExecCount.Load())
- assert.Equal(t, int64(0), sc2.ExecCount.Load())
+ verifyExpectedCount(t, sc1, 1)
+ verifyExpectedCount(t, sc2, 0)
 
  // no failure
  hc.Reset()
  sc1 = hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
  sc2 = hc.AddTestTablet("cell2", host, port, keyspace, shard, tabletType, true, 10, nil)
  err = f(ctx, tg, target)
  assert.NoError(t, err)
- assert.Equal(t, int64(0), sc1.ExecCount.Load())
- assert.Equal(t, int64(1), sc2.ExecCount.Load())
+ verifyExpectedCount(t, sc1, 0)
+ verifyExpectedCount(t, sc2, 1)
 
  // retry successful to other cell
  hc.Reset()
@@ -271,8 +285,8 @@ func testTabletGatewayGenericHelper(t *testing.T, ctx context.Context, f func(ct
  sc1.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
  err = f(ctx, tg, target)
  assert.NoError(t, err)
- assert.Equal(t, int64(1), sc1.ExecCount.Load())
- assert.Equal(t, int64(1), sc2.ExecCount.Load())
+ verifyExpectedCount(t, sc1, 1)
+ verifyExpectedCount(t, sc2, 1)
 }
 
 func testTabletGatewayTransact(t *testing.T, ctx context.Context, f func(ctx context.Context, tg *TabletGateway, target *querypb.Target) error) {