From 93b512be47b89b8acf7c6e49e79e6936b41a066c Mon Sep 17 00:00:00 2001
From: Jinzhu <wosmvp@gmail.com>
Date: Fri, 25 Dec 2015 19:33:57 +0800
Subject: [PATCH] Refactor association Clear, Replace

---
 association.go      | 34 ++++------------------------------
 association_test.go | 29 ++++++++++++++++++++++++++---
 2 files changed, 30 insertions(+), 33 deletions(-)

diff --git a/association.go b/association.go
index 39396e16..7f56c2d8 100644
--- a/association.go
+++ b/association.go
@@ -113,7 +113,7 @@ func (association *Association) Replace(values ...interface{}) *Association {
 	var foreignKeyMap = map[string]interface{}{}
 	for idx, foreignKey := range relationship.ForeignDBNames {
 		foreignKeyMap[foreignKey] = nil
-		if field, ok := scope.FieldByName(relationship.ForeignFieldNames[idx]); ok {
+		if field, ok := scope.FieldByName(relationship.AssociationForeignFieldNames[idx]); ok {
 			query = query.Where(fmt.Sprintf("%v = ?", scope.Quote(foreignKey)), field.Field.Interface())
 		}
 	}
@@ -125,6 +125,8 @@ func (association *Association) Replace(values ...interface{}) *Association {
 
 	if relationship.Kind == "many_to_many" {
 		association.setErr(relationship.JoinTableHandler.Delete(relationship.JoinTableHandler, query, relationship))
+	} else if relationship.Kind == "belongs_to" {
+		association.setErr(query.Model(scope.Value).UpdateColumn(foreignKeyMap).Error)
 	} else if relationship.Kind == "has_one" || relationship.Kind == "has_many" {
 		fieldValue := reflect.New(association.Field.Field.Type()).Interface()
 		association.setErr(query.Model(fieldValue).UpdateColumn(foreignKeyMap).Error)
@@ -210,35 +212,7 @@ func (association *Association) Delete(values ...interface{}) *Association {
 }
 
 func (association *Association) Clear() *Association {
-	relationship := association.Field.Relationship
-	scope := association.Scope
-	query := scope.NewDB()
-
-	if relationship.Kind == "many_to_many" {
-		for idx, foreignKey := range relationship.ForeignDBNames {
-			if field, ok := scope.FieldByName(relationship.ForeignFieldNames[idx]); ok {
-				query = query.Where(fmt.Sprintf("%v = ?", scope.Quote(foreignKey)), field.Field.Interface())
-			}
-		}
-
-		if err := relationship.JoinTableHandler.Delete(relationship.JoinTableHandler, query, relationship); err == nil {
-			association.Field.Set(reflect.Zero(association.Field.Field.Type()))
-		} else {
-			association.setErr(err)
-		}
-	} else {
-		association.Field.Set(reflect.Zero(association.Field.Field.Type()))
-		if relationship.Kind == "belongs_to" {
-			var foreignKeyMap = map[string]interface{}{}
-			for _, foreignKey := range relationship.ForeignDBNames {
-				foreignKeyMap[foreignKey] = nil
-			}
-			query.Model(scope.Value).Update(foreignKeyMap)
-		} else if relationship.Kind == "has_one" || relationship.Kind == "has_many" {
-			// query.Model(association.Field).UpdateColumn(foreignKeyMap)
-		}
-	}
-	return association
+	return association.Replace()
 }
 
 func (association *Association) Count() int {
diff --git a/association_test.go b/association_test.go
index b2d091c2..c094b81e 100644
--- a/association_test.go
+++ b/association_test.go
@@ -106,7 +106,7 @@ func TestBelongsTo(t *testing.T) {
 		t.Errorf("Should find category after append")
 	}
 
-	DB.Model(&post).Association("Category").Clear()
+	DB.Model(&post).Debug().Association("Category").Clear()
 
 	if !DB.Model(&post).Related(&Category{}).RecordNotFound() {
 		t.Errorf("Should not find any category after Clear")
@@ -179,15 +179,38 @@ func TestHasMany(t *testing.T) {
 	}
 
 	// Replace
-	DB.Model(&post).Association("Comments").Replace(&Comment{Content: "Comment 4"}, &Comment{Content: "Comment 5"})
+	DB.Model(&Post{Id: 999}).Debug().Association("Comments").Replace()
 
 	var comments4 []Comment
 	DB.Model(&post).Related(&comments4)
-	if !compareComments(comments4, []string{"Comment 4", "Comment 5"}) {
+	if len(comments4) == 0 {
+		t.Errorf("Replace should not clear all comments")
+	}
+
+	DB.Model(&post).Association("Comments").Replace(&Comment{Content: "Comment 4"}, &Comment{Content: "Comment 5"})
+
+	var comments41 []Comment
+	DB.Model(&post).Related(&comments41)
+	if !compareComments(comments41, []string{"Comment 4", "Comment 5"}) {
 		t.Errorf("Replace has many relations")
 	}
 
 	// Clear
+	DB.Model(&Post{Id: 999}).Association("Comments").Clear()
+
+	var comments5 []Comment
+	DB.Model(&post).Related(&comments5)
+	if len(comments5) == 0 {
+		t.Errorf("Clear should not clear all comments")
+	}
+
+	DB.Model(&post).Association("Comments").Clear()
+
+	var comments51 []Comment
+	DB.Model(&post).Related(&comments51)
+	if len(comments51) != 0 {
+		t.Errorf("Clear has many relations")
+	}
 }
 
 func TestHasOneAndHasManyAssociation(t *testing.T) {