Add nodeLabeller interface for testing

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

pkg/mig/reconfigure/api.go

@@ -30,3 +30,11 @@ type migParted interface {
 	applyMIGModeOnly() error
 	applyMIGConfig() error
 }
+
+// nodeLabeller defines an interface for interacting with node labels.
+//
+//go:generate moq -rm -fmt=goimports -out node-labeller_mock.go . nodeLabeller
+type nodeLabeller interface {
+	getNodeLabelValue(string) (string, error)
+	setNodeLabelValue(string, string) error
+}

pkg/mig/reconfigure/node-labeller_mock.go

@@ -29,6 +29,7 @@ import (
 
 	log "github.com/sirupsen/logrus"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes"
 )
 
 const (
@@ -41,8 +42,9 @@ const (
 
 type reconfigurer struct {
 	*reconfigureMIGOptions
-	migParted migParted
 	commandRunner
+	migParted migParted
+	node      nodeLabeller
 }
 
 // A commandWithOutput runs a command and ensures that STDERR and STDOUT are
@@ -74,6 +76,10 @@ func New(opts ...Option) (Reconfigurer, error) {
 			DriverLibraryPath:   o.DriverLibraryPath,
 			commandRunner:       c,
 		},
+		node: &node{
+			clientset: o.clientset,
+			name:      o.NodeName,
+		},
 	}
 
 	return r, nil
@@ -91,7 +97,7 @@ func (opts *reconfigurer) Reconfigure() error {
 	}
 
 	log.Infof("Getting current value of the '%s' node label", opts.ConfigStateLabel)
-	state, err := opts.getNodeLabelValue(opts.ConfigStateLabel)
+	state, err := opts.node.getNodeLabelValue(opts.ConfigStateLabel)
 	if err != nil {
 		return fmt.Errorf("unable to get the value of the %q label: %w", opts.ConfigStateLabel, err)
 	}
@@ -142,7 +148,7 @@ func (opts *reconfigurer) Reconfigure() error {
 	if err := opts.migParted.applyMIGModeOnly(); err != nil || opts.migParted.assertMIGModeOnly() != nil {
 		if opts.WithReboot {
 			log.Infof("Changing the '%s' node label to '%s'", opts.ConfigStateLabel, configStateRebooting)
-			if err := opts.setNodeLabelValue(opts.ConfigStateLabel, configStateRebooting); err != nil {
+			if err := opts.node.setNodeLabelValue(opts.ConfigStateLabel, configStateRebooting); err != nil {
 				log.Errorf("Unable to set the value of '%s' to '%s'", opts.ConfigStateLabel, configStateRebooting)
 				log.Error("Exiting so as not to reboot multiple times unexpectedly")
 				return fmt.Errorf("unable to set the value of %q to %q: %w", opts.ConfigStateLabel, configStateRebooting, err)
@@ -393,8 +399,13 @@ func rebootHost(hostRootMount string) error {
 	return nil
 }
 
-func (opts *reconfigureMIGOptions) getNodeLabelValue(label string) (string, error) {
-	node, err := opts.clientset.CoreV1().Nodes().Get(context.TODO(), opts.NodeName, metav1.GetOptions{})
+type node struct {
+	clientset *kubernetes.Clientset
+	name      string
+}
+
+func (n *node) getNodeLabelValue(label string) (string, error) {
+	node, err := n.clientset.CoreV1().Nodes().Get(context.TODO(), n.name, metav1.GetOptions{})
 	if err != nil {
 		return "", fmt.Errorf("unable to get node object: %w", err)
 	}
@@ -407,16 +418,16 @@ func (opts *reconfigureMIGOptions) getNodeLabelValue(label string) (string, erro
 	return value, nil
 }
 
-func (opts *reconfigureMIGOptions) setNodeLabelValue(label, value string) error {
-	node, err := opts.clientset.CoreV1().Nodes().Get(context.TODO(), opts.NodeName, metav1.GetOptions{})
+func (n *node) setNodeLabelValue(label, value string) error {
+	node, err := n.clientset.CoreV1().Nodes().Get(context.TODO(), n.name, metav1.GetOptions{})
 	if err != nil {
 		return fmt.Errorf("unable to get node object: %w", err)
 	}
 
 	labels := node.GetLabels()
 	labels[label] = value
 	node.SetLabels(labels)
-	_, err = opts.clientset.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
+	_, err = n.clientset.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
 	if err != nil {
 		return fmt.Errorf("unable to update node object: %w", err)
 	}

pkg/mig/reconfigure/reconfigure.go

@@ -35,15 +35,36 @@ func (c *commandRunnerWithCLI) Run(cmd *exec.Cmd) error {
 	return c.mock.Run(cmd)
 }
 
+type nodeWithLabels struct {
+	mock      *nodeLabellerMock
+	setLabels map[string]string
+}
+
+func (n *nodeWithLabels) getNodeLabelValue(label string) (string, error) {
+	return n.mock.getNodeLabelValue(label)
+}
+
+func (n *nodeWithLabels) setNodeLabelValue(label string, value string) error {
+	if err := n.mock.setNodeLabelValue(label, value); err != nil {
+		return err
+	}
+	if n.setLabels == nil {
+		n.setLabels = make(map[string]string)
+	}
+	n.setLabels[label] = value
+	return nil
+}
+
 func TestReconfigure(t *testing.T) {
 	testCases := []struct {
-		description    string
-		options        reconfigureMIGOptions
-		commandRunner  *commandRunnerWithCLI
-		migParted      *migPartedMock
-		checkMigParted func(*migPartedMock)
-		expectedError  error
-		expectedCalls  [][]string
+		description       string
+		options           reconfigureMIGOptions
+		migParted         *migPartedMock
+		checkMigParted    func(*migPartedMock)
+		nodeLabeller      *nodeWithLabels
+		checkNodeLabeller func(*nodeWithLabels)
+		expectedError     error
+		expectedCalls     [][]string
 	}{
 		{
 			description: "mig assert valid config failure does not call commands",
@@ -53,13 +74,7 @@ func TestReconfigure(t *testing.T) {
 				SelectedMIGConfig:   "selected-mig-config",
 				DriverLibraryPath:   "/path/to/libnvidia-ml.so.1",
 				HostRootMount:       "/host/",
-			},
-			commandRunner: &commandRunnerWithCLI{
-				mock: &commandRunnerMock{
-					RunFunc: func(cmd *exec.Cmd) error {
-						return fmt.Errorf("error running command %v", cmd.Path)
-					},
-				},
+				ConfigStateLabel:    "example.com/config.state",
 			},
 			migParted: &migPartedMock{
 				assertValidMIGConfigFunc: func() error {
@@ -76,9 +91,53 @@ func TestReconfigure(t *testing.T) {
 			expectedError: fmt.Errorf("error validating the selected MIG configuration: invalid mig config"),
 			expectedCalls: nil,
 		},
+		{
+			description: "node label error is causes exit",
+			options: reconfigureMIGOptions{
+				NodeName:            "NodeName",
+				MIGPartedConfigFile: "/path/to/config/file.yaml",
+				SelectedMIGConfig:   "selected-mig-config",
+				DriverLibraryPath:   "/path/to/libnvidia-ml.so.1",
+				HostRootMount:       "/host/",
+				ConfigStateLabel:    "example.com/config.state",
+			},
+			migParted: &migPartedMock{
+				assertValidMIGConfigFunc: func() error {
+					return nil
+				},
+			},
+			checkMigParted: func(mpm *migPartedMock) {
+				require.Len(t, mpm.calls.assertValidMIGConfig, 1)
+				require.Len(t, mpm.calls.applyMIGConfig, 0)
+				require.Len(t, mpm.calls.assertMIGModeOnly, 0)
+				require.Len(t, mpm.calls.applyMIGModeOnly, 0)
+				require.Len(t, mpm.calls.applyMIGConfig, 0)
+			},
+			nodeLabeller: &nodeWithLabels{
+				mock: &nodeLabellerMock{
+					getNodeLabelValueFunc: func(s string) (string, error) {
+						return "", fmt.Errorf("error getting label")
+					},
+				},
+			},
+			checkNodeLabeller: func(nwl *nodeWithLabels) {
+				calls := nwl.mock.getNodeLabelValueCalls()
+				require.Len(t, calls, 1)
+				require.EqualValues(t, []struct{ S string }{{"example.com/config.state"}}, calls)
+			},
+			expectedError: fmt.Errorf(`unable to get the value of the "example.com/config.state" label: error getting label`),
+		},
 	}
 
 	for _, tc := range testCases {
+		commandRunner := &commandRunnerWithCLI{
+			mock: &commandRunnerMock{
+				RunFunc: func(cmd *exec.Cmd) error {
+					return fmt.Errorf("error running command %v", cmd.Path)
+				},
+			},
+		}
+
 		t.Run(tc.description, func(t *testing.T) {
 			// TODO: Once we have better mocks in place for the following
 			// functionality, we can update this.
@@ -91,17 +150,26 @@ func TestReconfigure(t *testing.T) {
 
 			r := &reconfigurer{
 				reconfigureMIGOptions: &tc.options,
-				commandRunner:         tc.commandRunner,
+				commandRunner:         commandRunner,
 				migParted:             tc.migParted,
+				node:                  tc.nodeLabeller,
 			}
 
 			err := r.Reconfigure()
-			require.EqualValues(t, tc.expectedError.Error(), err.Error())
-
-			tc.checkMigParted(tc.migParted)
+			if tc.expectedError == nil {
+				require.NoError(t, err)
+			} else {
+				require.EqualError(t, err, tc.expectedError.Error())
+			}
 
-			require.EqualValues(t, tc.expectedCalls, tc.commandRunner.calls)
+			if tc.checkMigParted != nil {
+				tc.checkMigParted(tc.migParted)
+			}
+			if tc.checkNodeLabeller != nil {
+				tc.checkNodeLabeller(tc.nodeLabeller)
+			}
 
+			require.EqualValues(t, tc.expectedCalls, commandRunner.calls)
 		})
 	}
 }