From 0037c31a958ebf65a7d590a44cedc0671924f16d Mon Sep 17 00:00:00 2001
From: Xian Chaobo <xianchaobo@huawei.com>
Date: Tue, 7 Jul 2015 11:30:26 +0800
Subject: [PATCH] fix func name and use extracExprs

Signed-off-by: Xian Chaobo <xianchaobo@huawei.com>
---
 cluster/config.go        | 11 ++++-------
 cluster/config_test.go   |  4 ++--
 cluster/swarm/cluster.go |  2 +-
 3 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/cluster/config.go b/cluster/config.go
index 00eea8bad7..a46b5f5731 100644
--- a/cluster/config.go
+++ b/cluster/config.go
@@ -147,13 +147,10 @@ func (c *ContainerConfig) Constraints() []string {
 	return c.extractExprs("constraints")
 }
 
-// AddSoftImageAffinity to config
-func (c *ContainerConfig) AddSoftImageAffinity(image string) error {
-	affnitiesJSON := c.Labels[SwarmLabelNamespace+".affinities"]
-	var affinities []string
-	json.Unmarshal([]byte(affnitiesJSON), &affinities)
-	affinities = append(affinities, "image==~"+image)
-
+// AddAffinity to config
+func (c *ContainerConfig) AddAffinity(affinity string) error {
+	affinities := c.extractExprs("affinities")
+	affinities = append(affinities, affinity)
 	labels, err := json.Marshal(affinities)
 	if err != nil {
 		return err
diff --git a/cluster/config_test.go b/cluster/config_test.go
index faa0d7ef0f..15e0a6fe4f 100644
--- a/cluster/config_test.go
+++ b/cluster/config_test.go
@@ -82,10 +82,10 @@ func TestConsolidateResourceFields(t *testing.T) {
 
 }
 
-func TestAddSoftImageAffinity(t *testing.T) {
+func TestAddAffinity(t *testing.T) {
 	config := BuildContainerConfig(dockerclient.ContainerConfig{})
 	assert.Empty(t, config.Affinities())
 
-	config.AddSoftImageAffinity("testimage")
+	config.AddAffinity("image==~testimage")
 	assert.Len(t, config.Affinities(), 1)
 }
diff --git a/cluster/swarm/cluster.go b/cluster/swarm/cluster.go
index b756303e33..d7d960f2a5 100644
--- a/cluster/swarm/cluster.go
+++ b/cluster/swarm/cluster.go
@@ -116,7 +116,7 @@ func (c *Cluster) createContainer(config *cluster.ContainerConfig, name string,
 
 	configTemp := config
 	if withSoftImageAffinity {
-		configTemp.AddSoftImageAffinity(config.Image)
+		configTemp.AddAffinity("image==~" + config.Image)
 	}
 
 	n, err := c.scheduler.SelectNodeForContainer(c.listNodes(), configTemp)