From eb8dd093fe82f36e3bb548306f55e164bb3bd911 Mon Sep 17 00:00:00 2001
From: caiweidong <caiweidong@baidu.com>
Date: Tue, 19 Mar 2019 20:36:38 +0800
Subject: [PATCH] add UT

---
 .../statefulset/stateful_set_utils.go         |  4 +-
 .../statefulset/stateful_set_utils_test.go    | 63 +++++++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/pkg/controller/statefulset/stateful_set_utils.go b/pkg/controller/statefulset/stateful_set_utils.go
index a1cc09d07e..3c470fac43 100644
--- a/pkg/controller/statefulset/stateful_set_utils.go
+++ b/pkg/controller/statefulset/stateful_set_utils.go
@@ -151,9 +151,7 @@ func getPersistentVolumeClaims(set *apps.StatefulSet, pod *v1.Pod) map[string]v1
 		claim.Namespace = set.Namespace
 		if claim.Labels != nil {
 			for key, value := range set.Spec.Selector.MatchLabels {
-				if claim.Labels != nil {
-					claim.Labels[key] = value
-				}
+				claim.Labels[key] = value
 			}
 		} else {
 			claim.Labels = set.Spec.Selector.MatchLabels
diff --git a/pkg/controller/statefulset/stateful_set_utils_test.go b/pkg/controller/statefulset/stateful_set_utils_test.go
index cba555d9cb..ab4e9f04c6 100644
--- a/pkg/controller/statefulset/stateful_set_utils_test.go
+++ b/pkg/controller/statefulset/stateful_set_utils_test.go
@@ -19,6 +19,7 @@ package statefulset
 import (
 	"fmt"
 	"math/rand"
+	"reflect"
 	"sort"
 	"strconv"
 	"testing"
@@ -288,6 +289,68 @@ func TestCreateApplyRevision(t *testing.T) {
 	}
 }
 
+func TestGetPersistentVolumeClaims(t *testing.T) {
+
+	// nil inherits statefulset labels
+	pod := newPod()
+	statefulSet := newStatefulSet(1)
+	statefulSet.Spec.Selector.MatchLabels = nil
+	claims := getPersistentVolumeClaims(statefulSet, pod)
+	pvc := newPVC("datadir-foo-0")
+	pvc.SetNamespace(v1.NamespaceDefault)
+	resultClaims := map[string]v1.PersistentVolumeClaim{"datadir": pvc}
+
+	if !reflect.DeepEqual(claims, resultClaims) {
+		t.Fatalf("Unexpected pvc:\n %+v\n, desired pvc:\n %+v", claims, resultClaims)
+	}
+
+	// nil inherits statefulset labels
+	statefulSet.Spec.Selector.MatchLabels = map[string]string{"test": "test"}
+	claims = getPersistentVolumeClaims(statefulSet, pod)
+	pvc.SetLabels(map[string]string{"test": "test"})
+	resultClaims = map[string]v1.PersistentVolumeClaim{"datadir": pvc}
+	if !reflect.DeepEqual(claims, resultClaims) {
+		t.Fatalf("Unexpected pvc:\n %+v\n, desired pvc:\n %+v", claims, resultClaims)
+	}
+
+	// non-nil with non-overlapping labels merge pvc and statefulset labels
+	statefulSet.Spec.Selector.MatchLabels = map[string]string{"name": "foo"}
+	statefulSet.Spec.VolumeClaimTemplates[0].ObjectMeta.Labels = map[string]string{"test": "test"}
+	claims = getPersistentVolumeClaims(statefulSet, pod)
+	pvc.SetLabels(map[string]string{"test": "test", "name": "foo"})
+	resultClaims = map[string]v1.PersistentVolumeClaim{"datadir": pvc}
+	if !reflect.DeepEqual(claims, resultClaims) {
+		t.Fatalf("Unexpected pvc:\n %+v\n, desired pvc:\n %+v", claims, resultClaims)
+	}
+
+	// non-nil with overlapping labels merge pvc and statefulset labels and prefer statefulset labels
+	statefulSet.Spec.Selector.MatchLabels = map[string]string{"test": "foo"}
+	statefulSet.Spec.VolumeClaimTemplates[0].ObjectMeta.Labels = map[string]string{"test": "test"}
+	claims = getPersistentVolumeClaims(statefulSet, pod)
+	pvc.SetLabels(map[string]string{"test": "foo"})
+	resultClaims = map[string]v1.PersistentVolumeClaim{"datadir": pvc}
+	if !reflect.DeepEqual(claims, resultClaims) {
+		t.Fatalf("Unexpected pvc:\n %+v\n, desired pvc:\n %+v", claims, resultClaims)
+	}
+}
+
+func newPod() *v1.Pod {
+	return &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "foo-0",
+			Namespace: v1.NamespaceDefault,
+		},
+		Spec: v1.PodSpec{
+			Containers: []v1.Container{
+				{
+					Name:  "nginx",
+					Image: "nginx",
+				},
+			},
+		},
+	}
+}
+
 func newPVC(name string) v1.PersistentVolumeClaim {
 	return v1.PersistentVolumeClaim{
 		ObjectMeta: metav1.ObjectMeta{