From cfde8128f49e69be31ce433d7c9e298952356704 Mon Sep 17 00:00:00 2001
From: Harshit Gangal <harshit@planetscale.com>
Date: Thu, 19 Sep 2024 15:52:12 +0530
Subject: [PATCH] test: added check received error test

Signed-off-by: Harshit Gangal <harshit@planetscale.com>
---
 go/vt/vttablet/tabletserver/tx_engine_test.go | 58 +++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/go/vt/vttablet/tabletserver/tx_engine_test.go b/go/vt/vttablet/tabletserver/tx_engine_test.go
index 95057d754fb..be840c2a683 100644
--- a/go/vt/vttablet/tabletserver/tx_engine_test.go
+++ b/go/vt/vttablet/tabletserver/tx_engine_test.go
@@ -25,7 +25,10 @@ import (
 	"testing"
 	"time"
 
+	"vitess.io/vitess/go/mysql/sqlerror"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 	"vitess.io/vitess/go/vt/vtenv"
+	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tx"
 
 	"github.com/stretchr/testify/assert"
@@ -603,3 +606,58 @@ func TestTxEngineFailReserve(t *testing.T) {
 	require.Error(t, err)
 	assert.Zero(t, connID)
 }
+
+func TestCheckReceivedError(t *testing.T) {
+	db := setUpQueryExecutorTest(t)
+	defer db.Close()
+	db.AddQueryPattern(".*", &sqltypes.Result{})
+	cfg := tabletenv.NewDefaultConfig()
+	cfg.DB = newDBConfigs(db)
+	env := tabletenv.NewEnv(vtenv.NewTestEnv(), cfg, "TabletServerTest")
+	env.Config().TwoPCEnable = true
+	env.Config().TwoPCAbandonAge = 5
+	te := NewTxEngine(env, nil)
+	te.AcceptReadWrite()
+
+	tcases := []struct {
+		receivedErr  error
+		nonRetryable bool
+	}{{
+		receivedErr:  vterrors.New(vtrpcpb.Code_DEADLINE_EXCEEDED, "deadline exceeded"),
+		nonRetryable: false,
+	}, {
+		receivedErr:  vterrors.New(vtrpcpb.Code_INVALID_ARGUMENT, "invalid argument"),
+		nonRetryable: true,
+	}, {
+		receivedErr:  sqlerror.NewSQLError(sqlerror.ERLockDeadlock, sqlerror.SSLockDeadlock, "Deadlock found when trying to get lock; try restarting transaction"),
+		nonRetryable: true,
+	}, {
+		receivedErr:  context.DeadlineExceeded,
+		nonRetryable: false,
+	}, {
+		receivedErr:  context.Canceled,
+		nonRetryable: false,
+	}, {
+		receivedErr:  sqlerror.NewSQLError(sqlerror.CRServerLost, sqlerror.SSUnknownSQLState, "Lost connection to MySQL server during query"),
+		nonRetryable: false,
+	}, {
+		receivedErr:  sqlerror.NewSQLError(sqlerror.CRMalformedPacket, sqlerror.SSUnknownSQLState, "Malformed packet"),
+		nonRetryable: true,
+	}, {
+		receivedErr:  sqlerror.NewSQLError(sqlerror.CRServerGone, sqlerror.SSUnknownSQLState, "Server has gone away"),
+		nonRetryable: false,
+	}, {
+		receivedErr:  vterrors.New(vtrpcpb.Code_ABORTED, "Row count exceeded"),
+		nonRetryable: true,
+	}, {
+		receivedErr:  errors.New("(errno 2013) (sqlstate HY000) lost connection"),
+		nonRetryable: false,
+	}}
+
+	for _, tc := range tcases {
+		t.Run(tc.receivedErr.Error(), func(t *testing.T) {
+			nonRetryable := te.checkErrorAndMarkFailed(context.Background(), "aa", tc.receivedErr)
+			require.Equal(t, tc.nonRetryable, nonRetryable)
+		})
+	}
+}