@@ -40,3 +40,193 @@ func TestNewPartitionMultiUpdateFrom(t *testing.T) {
4040 require .Equal (t , ps .raw .Engine , op .(* PartitionMultiUpdate ).raw .Engine )
4141 require .Equal (t , ps .tableID , op .(* PartitionMultiUpdate ).tableID )
4242}
43+
44+ func TestAddInsertAffectRows (t * testing.T ) {
45+ tests := []struct {
46+ name string
47+ action actionType
48+ tableType UpdateTableType
49+ rowCount uint64
50+ expectRows uint64
51+ }{
52+ {
53+ name : "actionInsert with main table" ,
54+ action : actionInsert ,
55+ tableType : UpdateMainTable ,
56+ rowCount : 5 ,
57+ expectRows : 5 ,
58+ },
59+ {
60+ name : "actionUpdate with main table (REPLACE INTO)" ,
61+ action : actionUpdate ,
62+ tableType : UpdateMainTable ,
63+ rowCount : 3 ,
64+ expectRows : 3 ,
65+ },
66+ {
67+ name : "actionInsert with unique index table (should not count)" ,
68+ action : actionInsert ,
69+ tableType : UpdateUniqueIndexTable ,
70+ rowCount : 5 ,
71+ expectRows : 0 ,
72+ },
73+ {
74+ name : "actionUpdate with unique index table (should not count)" ,
75+ action : actionUpdate ,
76+ tableType : UpdateUniqueIndexTable ,
77+ rowCount : 3 ,
78+ expectRows : 0 ,
79+ },
80+ {
81+ name : "actionDelete with main table (should not count)" ,
82+ action : actionDelete ,
83+ tableType : UpdateMainTable ,
84+ rowCount : 2 ,
85+ expectRows : 0 ,
86+ },
87+ }
88+
89+ for _ , tt := range tests {
90+ t .Run (tt .name , func (t * testing.T ) {
91+ update := & MultiUpdate {
92+ ctr : container {
93+ action : tt .action ,
94+ affectedRows : 0 ,
95+ },
96+ }
97+ update .addAffectedRowsFunc = update .doAddAffectedRows
98+
99+ update .addInsertAffectRows (tt .tableType , tt .rowCount )
100+
101+ require .Equal (t , tt .expectRows , update .ctr .affectedRows , "affected rows should match expected value" )
102+ })
103+ }
104+ }
105+
106+ func TestAddDeleteAffectRows (t * testing.T ) {
107+ tests := []struct {
108+ name string
109+ action actionType
110+ tableType UpdateTableType
111+ rowCount uint64
112+ expectRows uint64
113+ }{
114+ {
115+ name : "actionDelete with main table" ,
116+ action : actionDelete ,
117+ tableType : UpdateMainTable ,
118+ rowCount : 5 ,
119+ expectRows : 5 ,
120+ },
121+ {
122+ name : "actionUpdate with main table (should not count for REPLACE INTO)" ,
123+ action : actionUpdate ,
124+ tableType : UpdateMainTable ,
125+ rowCount : 3 ,
126+ expectRows : 0 ,
127+ },
128+ {
129+ name : "actionDelete with unique index table (should not count)" ,
130+ action : actionDelete ,
131+ tableType : UpdateUniqueIndexTable ,
132+ rowCount : 5 ,
133+ expectRows : 0 ,
134+ },
135+ {
136+ name : "actionUpdate with unique index table (should not count)" ,
137+ action : actionUpdate ,
138+ tableType : UpdateUniqueIndexTable ,
139+ rowCount : 3 ,
140+ expectRows : 0 ,
141+ },
142+ {
143+ name : "actionInsert with main table (should not count)" ,
144+ action : actionInsert ,
145+ tableType : UpdateMainTable ,
146+ rowCount : 2 ,
147+ expectRows : 0 ,
148+ },
149+ }
150+
151+ for _ , tt := range tests {
152+ t .Run (tt .name , func (t * testing.T ) {
153+ update := & MultiUpdate {
154+ ctr : container {
155+ action : tt .action ,
156+ affectedRows : 0 ,
157+ },
158+ }
159+ update .addAffectedRowsFunc = update .doAddAffectedRows
160+
161+ update .addDeleteAffectRows (tt .tableType , tt .rowCount )
162+
163+ require .Equal (t , tt .expectRows , update .ctr .affectedRows , "affected rows should match expected value" )
164+ })
165+ }
166+ }
167+
168+ func TestReplaceIntoAffectedRows (t * testing.T ) {
169+ // Test REPLACE INTO scenario: should only count INSERT rows, not DELETE rows
170+ update := & MultiUpdate {
171+ ctr : container {
172+ action : actionUpdate , // REPLACE INTO uses actionUpdate
173+ affectedRows : 0 ,
174+ },
175+ }
176+ update .addAffectedRowsFunc = update .doAddAffectedRows
177+
178+ // Simulate REPLACE INTO: DELETE 2 rows, INSERT 2 rows
179+ // Should only count INSERT rows (2), not DELETE rows
180+ update .addDeleteAffectRows (UpdateMainTable , 2 ) // Should not count
181+ require .Equal (t , uint64 (0 ), update .ctr .affectedRows , "DELETE rows should not be counted for REPLACE INTO" )
182+
183+ update .addInsertAffectRows (UpdateMainTable , 2 ) // Should count
184+ require .Equal (t , uint64 (2 ), update .ctr .affectedRows , "INSERT rows should be counted for REPLACE INTO" )
185+ }
186+
187+ func TestUpdateAffectedRows (t * testing.T ) {
188+ // Test UPDATE scenario: should only count INSERT rows (updated rows), not DELETE rows
189+ update := & MultiUpdate {
190+ ctr : container {
191+ action : actionUpdate , // UPDATE uses actionUpdate
192+ affectedRows : 0 ,
193+ },
194+ }
195+ update .addAffectedRowsFunc = update .doAddAffectedRows
196+
197+ // Simulate UPDATE: DELETE 3 rows, INSERT 3 rows
198+ // Should only count INSERT rows (3), not DELETE rows
199+ update .addDeleteAffectRows (UpdateMainTable , 3 ) // Should not count
200+ require .Equal (t , uint64 (0 ), update .ctr .affectedRows , "DELETE rows should not be counted for UPDATE" )
201+
202+ update .addInsertAffectRows (UpdateMainTable , 3 ) // Should count
203+ require .Equal (t , uint64 (3 ), update .ctr .affectedRows , "INSERT rows should be counted for UPDATE" )
204+ }
205+
206+ func TestInsertAffectedRows (t * testing.T ) {
207+ // Test INSERT scenario: should count INSERT rows
208+ update := & MultiUpdate {
209+ ctr : container {
210+ action : actionInsert ,
211+ affectedRows : 0 ,
212+ },
213+ }
214+ update .addAffectedRowsFunc = update .doAddAffectedRows
215+
216+ update .addInsertAffectRows (UpdateMainTable , 5 )
217+ require .Equal (t , uint64 (5 ), update .ctr .affectedRows , "INSERT rows should be counted" )
218+ }
219+
220+ func TestDeleteAffectedRows (t * testing.T ) {
221+ // Test DELETE scenario: should count DELETE rows
222+ update := & MultiUpdate {
223+ ctr : container {
224+ action : actionDelete ,
225+ affectedRows : 0 ,
226+ },
227+ }
228+ update .addAffectedRowsFunc = update .doAddAffectedRows
229+
230+ update .addDeleteAffectRows (UpdateMainTable , 4 )
231+ require .Equal (t , uint64 (4 ), update .ctr .affectedRows , "DELETE rows should be counted" )
232+ }
0 commit comments