Fix FindInBatches tests

callbacks/create.go

@@ -55,6 +55,7 @@ func Create(config *Config) func(db *gorm.DB) {
 
 					if err == nil {
 						db.RowsAffected, _ = result.RowsAffected()
+
 						if db.RowsAffected > 0 {
 							if db.Statement.Schema != nil && db.Statement.Schema.PrioritizedPrimaryField != nil && db.Statement.Schema.PrioritizedPrimaryField.HasDefaultValue {
 								if insertID, err := result.LastInsertId(); err == nil && insertID > 0 {
@@ -138,6 +139,7 @@ func CreateWithReturning(db *gorm.DB) {
 			}
 
 			if !db.DryRun && db.Error == nil {
+				db.RowsAffected = 0
 				rows, err := db.Statement.ConnPool.QueryContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...)
 
 				if err == nil {

tests/query_test.go

@@ -260,13 +260,6 @@ func TestFindInBatches(t *testing.T) {
 	if result := DB.Where("name = ?", users[0].Name).FindInBatches(&results, 2, func(tx *gorm.DB, batch int) error {
 		totalBatch += batch
 
-		for idx := range results {
-			results[idx].Name = results[idx].Name + "_new"
-		}
-		if err := tx.Save(results).Error; err != nil {
-			t.Errorf("failed to save users, got error %v", err)
-		}
-
 		if tx.RowsAffected != 2 {
 			t.Errorf("Incorrect affected rows, expects: 2, got %v", tx.RowsAffected)
 		}
@@ -275,6 +268,14 @@ func TestFindInBatches(t *testing.T) {
 			t.Errorf("Incorrect users length, expects: 2, got %v", len(results))
 		}
 
+		for idx := range results {
+			results[idx].Name = results[idx].Name + "_new"
+		}
+
+		if err := tx.Save(results).Error; err != nil {
+			t.Errorf("failed to save users, got error %v", err)
+		}
+
 		return nil
 	}); result.Error != nil || result.RowsAffected != 6 {
 		t.Errorf("Failed to batch find, got error %v, rows affected: %v", result.Error, result.RowsAffected)