From 9d267e5285e48469e7408346e274366011b3c2fb Mon Sep 17 00:00:00 2001
From: Victor Vieux <victorvieux@gmail.com>
Date: Tue, 5 May 2015 17:24:19 -0700
Subject: [PATCH] fix affinities.bats  and add affinity_test.go

Signed-off-by: Victor Vieux <victorvieux@gmail.com>
---
 scheduler/filter/affinity_test.go | 60 +++++++++++++++++++++++++++++++
 test/integration/affinities.bats  | 17 +++++----
 2 files changed, 70 insertions(+), 7 deletions(-)

diff --git a/scheduler/filter/affinity_test.go b/scheduler/filter/affinity_test.go
index 91a21b90c9..f560b04baf 100644
--- a/scheduler/filter/affinity_test.go
+++ b/scheduler/filter/affinity_test.go
@@ -193,3 +193,63 @@ func TestAffinityFilter(t *testing.T) {
 	assert.Len(t, result, 0)
 
 }
+
+func TestAffinityFilterLabels(t *testing.T) {
+	var (
+		f     = AffinityFilter{}
+		nodes = []*node.Node{
+			{
+				ID:   "node-0-id",
+				Name: "node-0-name",
+				Addr: "node-0",
+				Containers: []*cluster.Container{
+					{Container: dockerclient.Container{
+						Id:    "container-n0-id",
+						Names: []string{"/container-n0-name"},
+					}},
+				},
+				Images: []*cluster.Image{{Image: dockerclient.Image{
+					Id:       "image-0-id",
+					RepoTags: []string{"image-0:tag0"},
+				}}},
+			},
+			{
+				ID:   "node-1-id",
+				Name: "node-1-name",
+				Addr: "node-1",
+				Containers: []*cluster.Container{
+					{Container: dockerclient.Container{
+						Id:    "container-n1-id",
+						Names: []string{"/container-n1-name"},
+					}},
+				},
+				Images: []*cluster.Image{{Image: dockerclient.Image{
+					Id:       "image-1-id",
+					RepoTags: []string{"image-1:tag1"},
+				}}},
+			},
+		}
+		result []*node.Node
+		err    error
+	)
+
+	result, err = f.Filter(cluster.BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"affinity:image==image-1"}}), nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 1)
+	assert.Equal(t, result[0], nodes[1])
+
+	result, err = f.Filter(cluster.BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"affinity:image!=image-1"}}), nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 1)
+	assert.Equal(t, result[0], nodes[0])
+
+	result, err = f.Filter(cluster.BuildContainerConfig(&dockerclient.ContainerConfig{Labels: map[string]string{"com.docker.swarm.affinities": "[\"image==image-1\"]"}}), nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 1)
+	assert.Equal(t, result[0], nodes[1])
+
+	result, err = f.Filter(cluster.BuildContainerConfig(&dockerclient.ContainerConfig{Labels: map[string]string{"com.docker.swarm.affinities": "[\"image!=image-1\"]"}}), nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 1)
+	assert.Equal(t, result[0], nodes[0])
+}
diff --git a/test/integration/affinities.bats b/test/integration/affinities.bats
index 0c797f5d25..b5fccff09a 100644
--- a/test/integration/affinities.bats
+++ b/test/integration/affinities.bats
@@ -44,20 +44,23 @@ function teardown() {
 }
 
 @test "image affinity" {
-	#FIXME: Broken
-	skip
 	start_docker 2
 	swarm_manage
 
-	run docker -H ${HOSTS[0]} pull busybox
+	run docker -H ${HOSTS[0]}  build -t test $BATS_TEST_DIRNAME/testdata/build
 	[ "$status" -eq 0 ]
-	run docker_swarm run --name c1 -e affinity:image==busybox -d busybox:latest sh
+
+	# pull busybox to force the refresh images
+	run docker_swarm pull busybox
 	[ "$status" -eq 0 ]
-	run docker_swarm run --name c2 -e affinity:image!=busybox -d busybox:latest sh
+
+	run docker_swarm run --name c1 -e affinity:image==test -d busybox:latest sh
 	[ "$status" -eq 0 ]
-	run docker_swarm run --name c3 --label 'com.docker.swarm.affinities=["image==busybox"]' -d busybox:latest sh
+	run docker_swarm run --name c2 -e affinity:image!=test -d busybox:latest sh
 	[ "$status" -eq 0 ]
-	run docker_swarm run --name c4 --label 'com.docker.swarm.affinities=["image\!=busybox"]' -d busybox:latest sh
+	run docker_swarm run --name c3 --label 'com.docker.swarm.affinities=["image==test"]' -d busybox:latest sh
+	[ "$status" -eq 0 ]
+	run docker_swarm run --name c4 --label 'com.docker.swarm.affinities=["image\!=test"]' -d busybox:latest sh
 	[ "$status" -eq 0 ]
 
 	run docker_swarm inspect c1