From d0b15030a3aa7a9d992c6c69d54202ea601b89a5 Mon Sep 17 00:00:00 2001
From: Cao Kai <ck89119@gmail.com>
Date: Tue, 17 Dec 2024 16:33:55 +0800
Subject: [PATCH] add ut

---
 pkg/frontend/session_test.go | 10 ++++++++++
 pkg/sql/plan/types_mock.go   |  8 ++++++++
 pkg/sql/plan/utils.go        |  4 ++--
 pkg/sql/plan/utils_test.go   | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/pkg/frontend/session_test.go b/pkg/frontend/session_test.go
index 60d24511628e0..60142350a19ba 100644
--- a/pkg/frontend/session_test.go
+++ b/pkg/frontend/session_test.go
@@ -384,6 +384,16 @@ func TestSession_TxnCompilerContext(t *testing.T) {
 		convey.So(object, convey.ShouldNotBeNil)
 		convey.So(tableRef, convey.ShouldNotBeNil)
 
+		ref := &plan.ObjectRef{
+			SchemaName: "schema",
+			PubInfo: &plan.PubInfo{
+				TenantId: 0,
+			},
+		}
+		object, tableRef = tcc.ResolveIndexTableByRef(ref, "indexTable", &plan2.Snapshot{TS: ts})
+		convey.So(object, convey.ShouldNotBeNil)
+		convey.So(tableRef, convey.ShouldNotBeNil)
+
 		pkd := tcc.GetPrimaryKeyDef("abc", "t1", &plan2.Snapshot{TS: ts})
 		convey.So(len(pkd), convey.ShouldBeZeroValue)
 
diff --git a/pkg/sql/plan/types_mock.go b/pkg/sql/plan/types_mock.go
index e062399894ca2..ed9accd307f2a 100644
--- a/pkg/sql/plan/types_mock.go
+++ b/pkg/sql/plan/types_mock.go
@@ -385,6 +385,14 @@ func (mr *MockCompilerContext2MockRecorder) ReplacePlan(execPlan interface{}) *g
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReplacePlan", reflect.TypeOf((*MockCompilerContext2)(nil).InitExecuteStmtParam), execPlan)
 }
 
+func (m *MockCompilerContext2) ResolveIndexTableByRef(ref *ObjectRef, tblName string, snapshot *Snapshot) (*ObjectRef, *TableDef) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ResolveIndexTableByRef", ref, tblName, snapshot)
+	ret0, _ := ret[0].(*ObjectRef)
+	ret1, _ := ret[1].(*TableDef)
+	return ret0, ret1
+}
+
 // Resolve mocks base method.
 func (m *MockCompilerContext2) Resolve(schemaName, tableName string, snapshot *Snapshot) (*ObjectRef, *TableDef) {
 	m.ctrl.T.Helper()
diff --git a/pkg/sql/plan/utils.go b/pkg/sql/plan/utils.go
index a6619bf7f9c33..cdb10f7a14feb 100644
--- a/pkg/sql/plan/utils.go
+++ b/pkg/sql/plan/utils.go
@@ -2721,10 +2721,10 @@ func offsetToString(offset int) string {
 // return !strings.HasPrefix(tableDef.Name, catalog.IndexTableNamePrefix)
 // }
 
-// DbNameOfObjRef return subscription name of ObjectRef if exists,
+// DbNameOfObjRef return subscription name of ObjectRef if exists, to avoid the mismatching of account id and db name
 func DbNameOfObjRef(objRef *ObjectRef) string {
 	if objRef.SubscriptionName == "" {
 		return objRef.SchemaName
 	}
 	return objRef.SubscriptionName
-}
\ No newline at end of file
+}
diff --git a/pkg/sql/plan/utils_test.go b/pkg/sql/plan/utils_test.go
index 2879e4c17cc8a..45d532946236b 100644
--- a/pkg/sql/plan/utils_test.go
+++ b/pkg/sql/plan/utils_test.go
@@ -101,3 +101,39 @@ func TestHandleOptimizerHints(t *testing.T) {
 	handleOptimizerHints("skipDedup=1", builder)
 	require.Equal(t, 1, builder.optimizerHints.skipDedup)
 }
+
+func TestDbNameOfObjRef(t *testing.T) {
+	type args struct {
+		objRef *ObjectRef
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "case 1",
+			args: args{
+				objRef: &ObjectRef{
+					SchemaName: "db",
+				},
+			},
+			want: "db",
+		},
+		{
+			name: "case 2",
+			args: args{
+				objRef: &ObjectRef{
+					SchemaName:       "whatever",
+					SubscriptionName: "sub",
+				},
+			},
+			want: "sub",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equalf(t, tt.want, DbNameOfObjRef(tt.args.objRef), "DbNameOfObjRef(%v)", tt.args.objRef)
+		})
+	}
+}