From ea782804cd6e0f589fa6805aa6df07ece7771c1e Mon Sep 17 00:00:00 2001
From: Manan Gupta <manan@planetscale.com>
Date: Fri, 20 Dec 2024 21:01:18 +0530
Subject: [PATCH] feat: add test for verifying dynamic config works

Signed-off-by: Manan Gupta <manan@planetscale.com>
---
 go/test/endtoend/cluster/vtgate_process.go    | 58 +++++++++++++++++++
 .../endtoend/transaction/twopc/main_test.go   |  3 +
 .../endtoend/transaction/twopc/twopc_test.go  | 32 ++++++++++
 3 files changed, 93 insertions(+)

diff --git a/go/test/endtoend/cluster/vtgate_process.go b/go/test/endtoend/cluster/vtgate_process.go
index 3e8c2635873..4253fbb5860 100644
--- a/go/test/endtoend/cluster/vtgate_process.go
+++ b/go/test/endtoend/cluster/vtgate_process.go
@@ -28,6 +28,7 @@ import (
 	"strconv"
 	"strings"
 	"syscall"
+	"testing"
 	"time"
 
 	"vitess.io/vitess/go/vt/log"
@@ -82,6 +83,63 @@ func (vtgate *VtgateProcess) RewriteConfiguration() error {
 	return os.WriteFile(vtgate.ConfigFile, []byte(vtgate.Config.ToJSONString()), 0644)
 }
 
+// WaitForConfig waits for the expectedConfig to be present in the vtgate configuration.
+func (vtgate *VtgateProcess) WaitForConfig(expectedConfig string) error {
+	timeout := time.After(30 * time.Second)
+	var response string
+	for {
+		select {
+		case <-timeout:
+			return fmt.Errorf("timed out waiting for api to work. Last response - %s", response)
+		default:
+			_, response, _ = vtgate.MakeAPICall("/debug/config")
+			if strings.Contains(response, expectedConfig) {
+				return nil
+			}
+			time.Sleep(1 * time.Second)
+		}
+	}
+}
+
+// MakeAPICall makes an API call on the given endpoint of VTOrc
+func (vtgate *VtgateProcess) MakeAPICall(endpoint string) (status int, response string, err error) {
+	url := fmt.Sprintf("http://localhost:%d/%s", vtgate.Port, endpoint)
+	resp, err := http.Get(url)
+	if err != nil {
+		if resp != nil {
+			status = resp.StatusCode
+		}
+		return status, "", err
+	}
+	defer func() {
+		if resp != nil && resp.Body != nil {
+			resp.Body.Close()
+		}
+	}()
+
+	respByte, _ := io.ReadAll(resp.Body)
+	return resp.StatusCode, string(respByte), err
+}
+
+// MakeAPICallRetry is used to make an API call and retries until success
+func (vtgate *VtgateProcess) MakeAPICallRetry(t *testing.T, url string) {
+	t.Helper()
+	timeout := time.After(10 * time.Second)
+	for {
+		select {
+		case <-timeout:
+			t.Fatal("timed out waiting for api to work")
+			return
+		default:
+			status, _, err := vtgate.MakeAPICall(url)
+			if err == nil && status == 200 {
+				return
+			}
+			time.Sleep(1 * time.Second)
+		}
+	}
+}
+
 const defaultVtGatePlannerVersion = planbuilder.Gen4
 
 // Setup starts Vtgate process with required arguements
diff --git a/go/test/endtoend/transaction/twopc/main_test.go b/go/test/endtoend/transaction/twopc/main_test.go
index 37337a62167..3607beea72a 100644
--- a/go/test/endtoend/transaction/twopc/main_test.go
+++ b/go/test/endtoend/transaction/twopc/main_test.go
@@ -106,6 +106,9 @@ func TestMain(m *testing.M) {
 		if err := clusterInstance.VtgateProcess.RewriteConfiguration(); err != nil {
 			return 1
 		}
+		if err := clusterInstance.VtgateProcess.WaitForConfig(`"transaction_mode":"TWOPC"`); err != nil {
+			return 1
+		}
 		vtParams = clusterInstance.GetVTParams(keyspaceName)
 		vtgateGrpcAddress = fmt.Sprintf("%s:%d", clusterInstance.Hostname, clusterInstance.VtgateGrpcPort)
 
diff --git a/go/test/endtoend/transaction/twopc/twopc_test.go b/go/test/endtoend/transaction/twopc/twopc_test.go
index a760cfb24b3..b7f7c11fba9 100644
--- a/go/test/endtoend/transaction/twopc/twopc_test.go
+++ b/go/test/endtoend/transaction/twopc/twopc_test.go
@@ -44,6 +44,38 @@ import (
 	"vitess.io/vitess/go/vt/vttablet/grpctmclient"
 )
 
+// TestDynamicConfig tests that transaction mode is dynamically configurable.
+func TestDynamicConfig(t *testing.T) {
+	conn, closer := start(t)
+	defer closer()
+	defer conn.Close()
+
+	// Ensure that initially running a distributed transaction is possible.
+	utils.Exec(t, conn, "begin")
+	utils.Exec(t, conn, "insert into twopc_t1(id, col) values(4, 4)")
+	utils.Exec(t, conn, "insert into twopc_t1(id, col) values(6, 4)")
+	utils.Exec(t, conn, "insert into twopc_t1(id, col) values(9, 4)")
+	utils.Exec(t, conn, "commit")
+
+	clusterInstance.VtgateProcess.Config.TransactionMode = "SINGLE"
+	defer func() {
+		clusterInstance.VtgateProcess.Config.TransactionMode = "TWOPC"
+		err := clusterInstance.VtgateProcess.RewriteConfiguration()
+		require.NoError(t, err)
+	}()
+	err := clusterInstance.VtgateProcess.RewriteConfiguration()
+	require.NoError(t, err)
+	err = clusterInstance.VtgateProcess.WaitForConfig(`"transaction_mode":"SINGLE"`)
+	require.NoError(t, err)
+
+	// After the config changes verify running a distributed transaction fails.
+	utils.Exec(t, conn, "begin")
+	utils.Exec(t, conn, "insert into twopc_t1(id, col) values(20, 4)")
+	_, err = utils.ExecAllowError(t, conn, "insert into twopc_t1(id, col) values(22, 4)")
+	require.ErrorContains(t, err, "multi-db transaction attempted")
+	utils.Exec(t, conn, "rollback")
+}
+
 // TestDTCommit tests distributed transaction commit for insert, update and delete operations
 // It verifies the binlog events for the same with transaction state changes and redo statements.
 func TestDTCommit(t *testing.T) {