From d63f5109a7553220f5d5d6cef58653e0e7e07435 Mon Sep 17 00:00:00 2001 From: Michael Demmer Date: Thu, 11 Jul 2024 07:51:43 -0700 Subject: [PATCH] add test helper to check the right count Signed-off-by: Michael Demmer --- go/vt/vtgate/tabletgateway_test.go | 52 +++++++++++++++++++----------- 1 file changed, 33 insertions(+), 19 deletions(-) diff --git a/go/vt/vtgate/tabletgateway_test.go b/go/vt/vtgate/tabletgateway_test.go index a867b09854d..b1e79b7803d 100644 --- a/go/vt/vtgate/tabletgateway_test.go +++ b/go/vt/vtgate/tabletgateway_test.go @@ -34,6 +34,7 @@ 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) { @@ -41,7 +42,10 @@ func TestTabletGatewayExecute(t *testing.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,7 +59,10 @@ 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) { @@ -63,7 +70,10 @@ func TestTabletGatewayBegin(t *testing.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,8 +266,8 @@ 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() @@ -261,8 +275,8 @@ func testTabletGatewayGenericHelper(t *testing.T, ctx context.Context, f func(ct 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) {