@@ -35,15 +35,36 @@ func (c *commandRunnerWithCLI) Run(cmd *exec.Cmd) error {
3535 return c .mock .Run (cmd )
3636}
3737
38+ type nodeWithLabels struct {
39+ mock * nodeLabellerMock
40+ setLabels map [string ]string
41+ }
42+
43+ func (n * nodeWithLabels ) getNodeLabelValue (label string ) (string , error ) {
44+ return n .mock .getNodeLabelValue (label )
45+ }
46+
47+ func (n * nodeWithLabels ) setNodeLabelValue (label string , value string ) error {
48+ if err := n .mock .setNodeLabelValue (label , value ); err != nil {
49+ return err
50+ }
51+ if n .setLabels == nil {
52+ n .setLabels = make (map [string ]string )
53+ }
54+ n .setLabels [label ] = value
55+ return nil
56+ }
57+
3858func TestReconfigure (t * testing.T ) {
3959 testCases := []struct {
40- description string
41- options reconfigureMIGOptions
42- commandRunner * commandRunnerWithCLI
43- migParted * migPartedMock
44- checkMigParted func (* migPartedMock )
45- expectedError error
46- expectedCalls [][]string
60+ description string
61+ options reconfigureMIGOptions
62+ migParted * migPartedMock
63+ checkMigParted func (* migPartedMock )
64+ nodeLabeller * nodeWithLabels
65+ checkNodeLabeller func (* nodeWithLabels )
66+ expectedError error
67+ expectedCalls [][]string
4768 }{
4869 {
4970 description : "mig assert valid config failure does not call commands" ,
@@ -53,13 +74,7 @@ func TestReconfigure(t *testing.T) {
5374 SelectedMIGConfig : "selected-mig-config" ,
5475 DriverLibraryPath : "/path/to/libnvidia-ml.so.1" ,
5576 HostRootMount : "/host/" ,
56- },
57- commandRunner : & commandRunnerWithCLI {
58- mock : & commandRunnerMock {
59- RunFunc : func (cmd * exec.Cmd ) error {
60- return fmt .Errorf ("error running command %v" , cmd .Path )
61- },
62- },
77+ ConfigStateLabel : "example.com/config.state" ,
6378 },
6479 migParted : & migPartedMock {
6580 assertValidMIGConfigFunc : func () error {
@@ -76,9 +91,53 @@ func TestReconfigure(t *testing.T) {
7691 expectedError : fmt .Errorf ("error validating the selected MIG configuration: invalid mig config" ),
7792 expectedCalls : nil ,
7893 },
94+ {
95+ description : "node label error is causes exit" ,
96+ options : reconfigureMIGOptions {
97+ NodeName : "NodeName" ,
98+ MIGPartedConfigFile : "/path/to/config/file.yaml" ,
99+ SelectedMIGConfig : "selected-mig-config" ,
100+ DriverLibraryPath : "/path/to/libnvidia-ml.so.1" ,
101+ HostRootMount : "/host/" ,
102+ ConfigStateLabel : "example.com/config.state" ,
103+ },
104+ migParted : & migPartedMock {
105+ assertValidMIGConfigFunc : func () error {
106+ return nil
107+ },
108+ },
109+ checkMigParted : func (mpm * migPartedMock ) {
110+ require .Len (t , mpm .calls .assertValidMIGConfig , 1 )
111+ require .Len (t , mpm .calls .applyMIGConfig , 0 )
112+ require .Len (t , mpm .calls .assertMIGModeOnly , 0 )
113+ require .Len (t , mpm .calls .applyMIGModeOnly , 0 )
114+ require .Len (t , mpm .calls .applyMIGConfig , 0 )
115+ },
116+ nodeLabeller : & nodeWithLabels {
117+ mock : & nodeLabellerMock {
118+ getNodeLabelValueFunc : func (s string ) (string , error ) {
119+ return "" , fmt .Errorf ("error getting label" )
120+ },
121+ },
122+ },
123+ checkNodeLabeller : func (nwl * nodeWithLabels ) {
124+ calls := nwl .mock .getNodeLabelValueCalls ()
125+ require .Len (t , calls , 1 )
126+ require .EqualValues (t , []struct { S string }{{"example.com/config.state" }}, calls )
127+ },
128+ expectedError : fmt .Errorf (`unable to get the value of the "example.com/config.state" label: error getting label` ),
129+ },
79130 }
80131
81132 for _ , tc := range testCases {
133+ commandRunner := & commandRunnerWithCLI {
134+ mock : & commandRunnerMock {
135+ RunFunc : func (cmd * exec.Cmd ) error {
136+ return fmt .Errorf ("error running command %v" , cmd .Path )
137+ },
138+ },
139+ }
140+
82141 t .Run (tc .description , func (t * testing.T ) {
83142 // TODO: Once we have better mocks in place for the following
84143 // functionality, we can update this.
@@ -91,17 +150,26 @@ func TestReconfigure(t *testing.T) {
91150
92151 r := & reconfigurer {
93152 reconfigureMIGOptions : & tc .options ,
94- commandRunner : tc . commandRunner ,
153+ commandRunner : commandRunner ,
95154 migParted : tc .migParted ,
155+ node : tc .nodeLabeller ,
96156 }
97157
98158 err := r .Reconfigure ()
99- require .EqualValues (t , tc .expectedError .Error (), err .Error ())
100-
101- tc .checkMigParted (tc .migParted )
159+ if tc .expectedError == nil {
160+ require .NoError (t , err )
161+ } else {
162+ require .EqualError (t , err , tc .expectedError .Error ())
163+ }
102164
103- require .EqualValues (t , tc .expectedCalls , tc .commandRunner .calls )
165+ if tc .checkMigParted != nil {
166+ tc .checkMigParted (tc .migParted )
167+ }
168+ if tc .checkNodeLabeller != nil {
169+ tc .checkNodeLabeller (tc .nodeLabeller )
170+ }
104171
172+ require .EqualValues (t , tc .expectedCalls , commandRunner .calls )
105173 })
106174 }
107175}
0 commit comments