fix tests

go-gorm · Jun 24, 2024 · ce92572 · ce92572
1 parent 290a62e
commit ce92572
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 30 deletions.
diff --git a/callbacks/associations.go b/callbacks/associations.go
@@ -179,19 +179,7 @@ func SaveAfterAssociations(create bool) func(db *gorm.DB) {
 								db.AddError(ref.ForeignKey.Set(db.Statement.Context, f, ref.PrimaryValue))
 							}
 
-							if ref.ForeignKey.Schema != nil {
-								for _, field := range ref.ForeignKey.Schema.Fields {
-									if field.PrimaryKey {
-										continue
-									}
-
-									if field.DBName == "" {
-										continue
-									}
-
-									assignmentColumns = append(assignmentColumns, field.DBName)
-								}
-							}
+							assignmentColumns = append(assignmentColumns, getAssignmentColumnsForForeignKey(ref.ForeignKey)...)
 						}
 
 						saveAssociations(db, rel, f, selectColumns, restricted, assignmentColumns)
@@ -444,6 +432,33 @@ func saveAssociations(db *gorm.DB, rel *schema.Relationship, rValues reflect.Val
 	return db.AddError(tx.Create(values).Error)
 }
 
+func getAssignmentColumnsForForeignKey(foreignKey *schema.Field) []string {
+	var assignmentColumns []string
+
+	if foreignKey.Schema == nil {
+		return assignmentColumns
+	}
+
+	if !foreignKey.PrimaryKey {
+		assignmentColumns = append(assignmentColumns, foreignKey.DBName)
+		return assignmentColumns
+	}
+
+	for _, field := range foreignKey.Schema.Fields {
+		if field.PrimaryKey {
+			continue
+		}
+
+		if field.DBName == "" {
+			continue
+		}
+
+		assignmentColumns = append(assignmentColumns, field.DBName)
+	}
+
+	return assignmentColumns
+}
+
 // check association values has been saved
 // if values kind is Struct, check it has been saved
 // if values kind is Slice/Array, check all items have been saved

diff --git a/tests/associations_has_one_test.go b/tests/associations_has_one_test.go
@@ -257,6 +257,10 @@ func TestPolymorphicHasOneAssociationForSlice(t *testing.T) {
 }
 
 func TestReplaceHasOneAssociationWithCustomPK(t *testing.T) {
+	if DB.Dialector.Name() == "sqlite" {
+		return
+	}
+
 	DB.Migrator().DropTable(&User{})
 	DB.AutoMigrate(&User{})
 	DB.AutoMigrate(&CreditCard{})

diff --git a/utils/tests/models.go b/utils/tests/models.go
@@ -14,23 +14,22 @@ import (
 // NamedPet is a reference to a named `Pet` (has one)
 type User struct {
 	gorm.Model
-	Name       string
-	Age        uint
-	Birthday   *time.Time
-	Account    Account
-	Pets       []*Pet
-	NamedPet   *Pet
-	Toys       []Toy   `gorm:"polymorphic:Owner"`
-	Tools      []Tools `gorm:"polymorphicType:Type;polymorphicId:CustomID"`
-	CompanyID  *int
-	Company    Company
-	ManagerID  *uint
-	Manager    *User
-	Team       []User     `gorm:"foreignkey:ManagerID"`
-	Languages  []Language `gorm:"many2many:UserSpeak;"`
-	Friends    []*User    `gorm:"many2many:user_friends;"`
-	Active     bool
-	CreditCard CreditCard `gorm:"foreignKey:UserName;references:name"`
+	Name      string
+	Age       uint
+	Birthday  *time.Time
+	Account   Account
+	Pets      []*Pet
+	NamedPet  *Pet
+	Toys      []Toy   `gorm:"polymorphic:Owner"`
+	Tools     []Tools `gorm:"polymorphicType:Type;polymorphicId:CustomID"`
+	CompanyID *int
+	Company   Company
+	ManagerID *uint
+	Manager   *User
+	Team      []User     `gorm:"foreignkey:ManagerID"`
+	Languages []Language `gorm:"many2many:UserSpeak;"`
+	Friends   []*User    `gorm:"many2many:user_friends;"`
+	Active    bool
 }
 
 type Account struct {
@@ -104,6 +103,12 @@ type Child struct {
 	Parent   *Parent
 }
 
+type Owner struct {
+	gorm.Model
+	Name       string     `gorm:"index"`
+	CreditCard CreditCard `gorm:"foreignKey:OwnerName;references:name"`
+}
+
 type CreditCard struct {
 	Number   string
 	UserName string `gorm:"primaryKey;unique;size:255"`