diff --git a/main.go b/main.go index 272b1d142..5850c3c3a 100644 --- a/main.go +++ b/main.go @@ -6,6 +6,7 @@ import ( "log" "os" "regexp" + "sort" "strings" "time" @@ -78,10 +79,13 @@ type APIHelpers interface { // GetNode returns the Kubernetes node on which this container is running. GetNode(*k8sclient.Clientset) (*api.Node, error) - // RemoveLabels removes labels from the supplied node that contain the + // RemoveLabelsWithPrefix removes labels from the supplied node that contain the // search string provided. In order to publish the changes, the node must // subsequently be updated via the API server using the client library. - RemoveLabels(*api.Node, string) + RemoveLabelsWithPrefix(*api.Node, string) + + // RemoveLabels removes NFD labels from a node object + RemoveLabels(*api.Node, []string) // AddLabels adds new NFD labels to the node object. // In order to publish the labels, the node must be subsequently updated via the @@ -314,8 +318,14 @@ func createFeatureLabels(sources []source.FeatureSource, labelWhiteList *regexp. // disabled via --no-publish flag. func updateNodeWithFeatureLabels(helper APIHelpers, noPublish bool, labels Labels) error { if !noPublish { - // Advertise NFD version as an annotation - annotations := Annotations{"version": version} + // Advertise NFD version and label names as annotations + keys := make([]string, 0, len(labels)) + for k, _ := range labels { + keys = append(keys, k) + } + sort.Strings(keys) + annotations := Annotations{"version": version, + "feature-labels": strings.Join(keys, ",")} err := advertiseFeatureLabels(helper, labels, annotations) if err != nil { @@ -363,11 +373,15 @@ func advertiseFeatureLabels(helper APIHelpers, labels Labels, annotations Annota return err } - // Remove labels with our prefix - helper.RemoveLabels(node, labelPrefix) + // Remove old labels + if l, ok := node.Annotations[annotationNs+"feature-labels"]; ok { + oldLabels := strings.Split(l, ",") + helper.RemoveLabels(node, oldLabels) + } + // Also, remove all labels with the old prefix, and the old version label - helper.RemoveLabels(node, "node.alpha.kubernetes-incubator.io/nfd") - helper.RemoveLabels(node, "node.alpha.kubernetes-incubator.io/node-feature-discovery") + helper.RemoveLabelsWithPrefix(node, "node.alpha.kubernetes-incubator.io/nfd") + helper.RemoveLabelsWithPrefix(node, "node.alpha.kubernetes-incubator.io/node-feature-discovery") // Add labels to the node object. helper.AddLabels(node, labels) @@ -416,9 +430,9 @@ func (h k8sHelpers) GetNode(cli *k8sclient.Clientset) (*api.Node, error) { return node, nil } -// RemoveLabels searches through all labels on Node n and removes +// RemoveLabelsWithPrefix searches through all labels on Node n and removes // any where the key contain the search string. -func (h k8sHelpers) RemoveLabels(n *api.Node, search string) { +func (h k8sHelpers) RemoveLabelsWithPrefix(n *api.Node, search string) { for k := range n.Labels { if strings.Contains(k, search) { delete(n.Labels, k) @@ -426,6 +440,13 @@ func (h k8sHelpers) RemoveLabels(n *api.Node, search string) { } } +// RemoveLabels removes given NFD labels +func (h k8sHelpers) RemoveLabels(n *api.Node, labelNames []string) { + for _, l := range labelNames { + delete(n.Labels, labelPrefix+l) + } +} + func (h k8sHelpers) AddLabels(n *api.Node, labels Labels) { for k, v := range labels { n.Labels[labelPrefix+k] = v diff --git a/main_test.go b/main_test.go index 158552456..618e618f6 100644 --- a/main_test.go +++ b/main_test.go @@ -5,6 +5,7 @@ import ( "io/ioutil" "os" "regexp" + "strings" "testing" "time" @@ -25,11 +26,16 @@ func TestDiscoveryWithMockSources(t *testing.T) { fakeFeatureNames := []string{"testfeature1", "testfeature2", "testfeature3"} fakeFeatures := source.Features{} fakeFeatureLabels := Labels{} - fakeAnnotations := Annotations{"version": version} + fakeAnnotations := Annotations{"version": version, + "feature-labels": "testSource-testfeature1,testSource-testfeature2,testSource-testfeature3"} + fakeFeatureLabelNames := make([]string, 0, len(fakeFeatureNames)) for _, f := range fakeFeatureNames { fakeFeatures[f] = true - fakeFeatureLabels[fmt.Sprintf("testSource-%s", f)] = "true" + labelName := fakeFeatureSourceName + "-" + f + fakeFeatureLabels[labelName] = "true" + fakeFeatureLabelNames = append(fakeFeatureLabelNames, labelName) } + fakeAnnotations["feature-labels"] = strings.Join(fakeFeatureLabelNames, ",") fakeFeatureSource := source.FeatureSource(mockFeatureSource) Convey("When I successfully get the labels from the mock source", func() { @@ -60,16 +66,16 @@ func TestDiscoveryWithMockSources(t *testing.T) { mockAPIHelper := new(MockAPIHelpers) testHelper := APIHelpers(mockAPIHelper) + mockNode := &api.Node{} var mockClient *k8sclient.Clientset - var mockNode *api.Node Convey("When I successfully update the node with feature labels", func() { mockAPIHelper.On("GetClient").Return(mockClient, nil) mockAPIHelper.On("GetNode", mockClient).Return(mockNode, nil).Once() mockAPIHelper.On("AddLabels", mockNode, fakeFeatureLabels).Return().Once() - mockAPIHelper.On("RemoveLabels", mockNode, labelPrefix).Return().Once() - mockAPIHelper.On("RemoveLabels", mockNode, "node.alpha.kubernetes-incubator.io/nfd").Return().Once() - mockAPIHelper.On("RemoveLabels", mockNode, "node.alpha.kubernetes-incubator.io/node-feature-discovery").Return().Once() + mockAPIHelper.On("RemoveLabelsWithPrefix", mockNode, labelPrefix).Return().Once() + mockAPIHelper.On("RemoveLabelsWithPrefix", mockNode, "node.alpha.kubernetes-incubator.io/nfd").Return().Once() + mockAPIHelper.On("RemoveLabelsWithPrefix", mockNode, "node.alpha.kubernetes-incubator.io/node-feature-discovery").Return().Once() mockAPIHelper.On("AddAnnotations", mockNode, fakeAnnotations).Return().Once() mockAPIHelper.On("UpdateNode", mockClient, mockNode).Return(nil).Once() noPublish := false @@ -116,9 +122,9 @@ func TestDiscoveryWithMockSources(t *testing.T) { expectedError := errors.New("fake error") mockAPIHelper.On("GetClient").Return(mockClient, nil) mockAPIHelper.On("GetNode", mockClient).Return(mockNode, nil).Once() - mockAPIHelper.On("RemoveLabels", mockNode, labelPrefix).Return().Once() - mockAPIHelper.On("RemoveLabels", mockNode, "node.alpha.kubernetes-incubator.io/nfd").Return().Once() - mockAPIHelper.On("RemoveLabels", mockNode, "node.alpha.kubernetes-incubator.io/node-feature-discovery").Return().Once() + mockAPIHelper.On("RemoveLabelsWithPrefix", mockNode, labelPrefix).Return().Once() + mockAPIHelper.On("RemoveLabelsWithPrefix", mockNode, "node.alpha.kubernetes-incubator.io/nfd").Return().Once() + mockAPIHelper.On("RemoveLabelsWithPrefix", mockNode, "node.alpha.kubernetes-incubator.io/node-feature-discovery").Return().Once() mockAPIHelper.On("AddLabels", mockNode, fakeFeatureLabels).Return().Once() mockAPIHelper.On("AddAnnotations", mockNode, fakeAnnotations).Return().Once() mockAPIHelper.On("UpdateNode", mockClient, mockNode).Return(expectedError).Once() @@ -338,7 +344,7 @@ func TestAddLabels(t *testing.T) { }) } -func TestRemoveLabels(t *testing.T) { +func TestRemoveLabelsWithPrefix(t *testing.T) { Convey("When removing labels", t, func() { helper := k8sHelpers{} n := &api.Node{ @@ -352,20 +358,20 @@ func TestRemoveLabels(t *testing.T) { } Convey("a unique label should be removed", func() { - helper.RemoveLabels(n, "single") + helper.RemoveLabelsWithPrefix(n, "single") So(len(n.Labels), ShouldEqual, 2) So(n.Labels, ShouldNotContainKey, "single") }) Convey("a non-unique search string should remove all matching keys", func() { - helper.RemoveLabels(n, "multiple") + helper.RemoveLabelsWithPrefix(n, "multiple") So(len(n.Labels), ShouldEqual, 1) So(n.Labels, ShouldNotContainKey, "multiple_A") So(n.Labels, ShouldNotContainKey, "multiple_B") }) Convey("a search string with no matches should not alter labels", func() { - helper.RemoveLabels(n, "unique") + helper.RemoveLabelsWithPrefix(n, "unique") So(n.Labels, ShouldContainKey, "single") So(n.Labels, ShouldContainKey, "multiple_A") So(n.Labels, ShouldContainKey, "multiple_B") diff --git a/mockapihelpers.go b/mockapihelpers.go index db55da91e..6bf5b18ff 100644 --- a/mockapihelpers.go +++ b/mockapihelpers.go @@ -58,9 +58,15 @@ func (_m *MockAPIHelpers) GetNode(_a0 *k8sclient.Clientset) (*api.Node, error) { return r0, r1 } -// RemoveLabels provides a mock function with *api.Node and main.Labels as the input arguments and +// RemoveLabelsWithPrefix provides a mock function with *api.Node and main.Labels as the input arguments and // no return value -func (_m *MockAPIHelpers) RemoveLabels(_a0 *api.Node, _a1 string) { +func (_m *MockAPIHelpers) RemoveLabelsWithPrefix(_a0 *api.Node, _a1 string) { + _m.Called(_a0, _a1) +} + +// RemoveLabels provides a mock function with *api.Node and []strings as the input arguments and +// no return value +func (_m *MockAPIHelpers) RemoveLabels(_a0 *api.Node, _a1 []string) { _m.Called(_a0, _a1) }