From b14f0e84106eeb9fdafa6f69eeb6309f81d8b210 Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Fri, 20 Dec 2024 15:37:55 +0100
Subject: [PATCH] test: refactor test

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 go/vt/vtgate/scatter_conn_test.go | 69 +++++++++++++++----------------
 1 file changed, 34 insertions(+), 35 deletions(-)

diff --git a/go/vt/vtgate/scatter_conn_test.go b/go/vt/vtgate/scatter_conn_test.go
index 2a4dddbd5b7..e5c27c0de33 100644
--- a/go/vt/vtgate/scatter_conn_test.go
+++ b/go/vt/vtgate/scatter_conn_test.go
@@ -108,6 +108,40 @@ func TestExecuteFailOnAutocommit(t *testing.T) {
 func TestFetchLastInsertIDResets(t *testing.T) {
 	// This test verifies that the FetchLastInsertID flag is reset after a call to ExecuteMultiShard.
 	ks := "TestFetchLastInsertIDResets"
+	ctx := utils.LeakCheckContext(t)
+
+	createSandbox(ks)
+	hc := discovery.NewFakeHealthCheck(nil)
+	sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa")
+	sbc0 := hc.AddTestTablet("aa", "0", 1, ks, "0", topodatapb.TabletType_PRIMARY, true, 1, nil)
+	sbc1 := hc.AddTestTablet("aa", "1", 1, ks, "1", topodatapb.TabletType_PRIMARY, true, 1, nil)
+
+	rss := []*srvtopo.ResolvedShard{{
+		Target: &querypb.Target{
+			Keyspace:   ks,
+			Shard:      "0",
+			TabletType: topodatapb.TabletType_PRIMARY,
+		},
+		Gateway: sbc0,
+	}, {
+		Target: &querypb.Target{
+			Keyspace:   ks,
+			Shard:      "1",
+			TabletType: topodatapb.TabletType_PRIMARY,
+		},
+		Gateway: sbc1,
+	}}
+	queries := []*querypb.BoundQuery{{
+		Sql: "query1",
+		BindVariables: map[string]*querypb.BindVariable{
+			"bv0": sqltypes.Int64BindVariable(0),
+		},
+	}, {
+		Sql: "query2",
+		BindVariables: map[string]*querypb.BindVariable{
+			"bv1": sqltypes.Int64BindVariable(1),
+		},
+	}}
 	tests := []struct {
 		name               string
 		initialSessionOpts *querypb.ExecuteOptions
@@ -148,41 +182,6 @@ func TestFetchLastInsertIDResets(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ctx := utils.LeakCheckContext(t)
-
-			createSandbox(ks)
-			hc := discovery.NewFakeHealthCheck(nil)
-			sc := newTestScatterConn(ctx, hc, newSandboxForCells(ctx, []string{"aa"}), "aa")
-			sbc0 := hc.AddTestTablet("aa", "0", 1, ks, "0", topodatapb.TabletType_PRIMARY, true, 1, nil)
-			sbc1 := hc.AddTestTablet("aa", "1", 1, ks, "1", topodatapb.TabletType_PRIMARY, true, 1, nil)
-
-			rss := []*srvtopo.ResolvedShard{{
-				Target: &querypb.Target{
-					Keyspace:   ks,
-					Shard:      "0",
-					TabletType: topodatapb.TabletType_PRIMARY,
-				},
-				Gateway: sbc0,
-			}, {
-				Target: &querypb.Target{
-					Keyspace:   ks,
-					Shard:      "1",
-					TabletType: topodatapb.TabletType_PRIMARY,
-				},
-				Gateway: sbc1,
-			}}
-			queries := []*querypb.BoundQuery{{
-				Sql: "query1",
-				BindVariables: map[string]*querypb.BindVariable{
-					"bv0": sqltypes.Int64BindVariable(0),
-				},
-			}, {
-				Sql: "query2",
-				BindVariables: map[string]*querypb.BindVariable{
-					"bv1": sqltypes.Int64BindVariable(1),
-				},
-			}}
-
 			session := econtext.NewSafeSession(nil)
 			session.Options = tt.initialSessionOpts