@@ -27,6 +27,7 @@ import (
2727 corev1 "k8s.io/api/core/v1"
2828 metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2929 "k8s.io/apimachinery/pkg/runtime"
30+ "k8s.io/klog/v2/klogr"
3031 "sigs.k8s.io/cluster-api-provider-azure/api/v1alpha4"
3132 "sigs.k8s.io/cluster-api-provider-azure/azure"
3233 mock_scope "sigs.k8s.io/cluster-api-provider-azure/azure/scope/mocks"
@@ -316,6 +317,100 @@ func TestMachineScope_UpdateStatus(t *testing.T) {
316317 }
317318}
318319
320+ func TestMachinePoolMachineScope_CordonAndDrain (t * testing.T ) {
321+ scheme := runtime .NewScheme ()
322+ _ = capiv1exp .AddToScheme (scheme )
323+ _ = infrav1 .AddToScheme (scheme )
324+
325+ var (
326+ clusterScope = ClusterScope {
327+ Cluster : & clusterv1.Cluster {
328+ ObjectMeta : metav1.ObjectMeta {
329+ Name : "cluster-foo" ,
330+ },
331+ },
332+ }
333+ )
334+
335+ cases := []struct {
336+ Name string
337+ Setup func (mockNodeGetter * mock_scope.MocknodeGetter , ampm * infrav1.AzureMachinePoolMachine ) * infrav1.AzureMachinePoolMachine
338+ Err string
339+ }{
340+ {
341+ Name : "should skip cordon and drain if the node does not exist with provider ID" ,
342+ Setup : func (mockNodeGetter * mock_scope.MocknodeGetter , ampm * infrav1.AzureMachinePoolMachine ) * infrav1.AzureMachinePoolMachine {
343+ mockNodeGetter .EXPECT ().GetNodeByProviderID (gomock2 .AContext (), FakeProviderID ).Return (nil , nil )
344+ return ampm
345+ },
346+ },
347+ {
348+ Name : "should skip cordon and drain if the node does not exist with node reference" ,
349+ Setup : func (mockNodeGetter * mock_scope.MocknodeGetter , ampm * infrav1.AzureMachinePoolMachine ) * infrav1.AzureMachinePoolMachine {
350+ nodeRef := corev1.ObjectReference {
351+ Name : "node1" ,
352+ }
353+ ampm .Status .NodeRef = & nodeRef
354+ mockNodeGetter .EXPECT ().GetNodeByObjectReference (gomock2 .AContext (), nodeRef ).Return (nil , nil )
355+ return ampm
356+ },
357+ },
358+ {
359+ Name : "if GetNodeByProviderID fails with an error, an error will be returned" ,
360+ Setup : func (mockNodeGetter * mock_scope.MocknodeGetter , ampm * infrav1.AzureMachinePoolMachine ) * infrav1.AzureMachinePoolMachine {
361+ mockNodeGetter .EXPECT ().GetNodeByProviderID (gomock2 .AContext (), FakeProviderID ).Return (nil , errors .New ("boom" ))
362+ return ampm
363+ },
364+ Err : "failed to find node: boom" ,
365+ },
366+ }
367+
368+ for _ , c := range cases {
369+ t .Run (c .Name , func (t * testing.T ) {
370+ var (
371+ controller = gomock .NewController (t )
372+ mockClient = mock_scope .NewMocknodeGetter (controller )
373+ g = NewWithT (t )
374+ params = MachinePoolMachineScopeParams {
375+ Client : fake .NewClientBuilder ().WithScheme (scheme ).Build (),
376+ ClusterScope : & clusterScope ,
377+ MachinePool : & capiv1exp.MachinePool {
378+ Spec : capiv1exp.MachinePoolSpec {
379+ Template : clusterv1.MachineTemplateSpec {
380+ Spec : clusterv1.MachineSpec {
381+ Version : to .StringPtr ("v1.19.11" ),
382+ },
383+ },
384+ },
385+ },
386+ AzureMachinePool : new (infrav1.AzureMachinePool ),
387+ }
388+ )
389+
390+ defer controller .Finish ()
391+
392+ ampm := c .Setup (mockClient , & infrav1.AzureMachinePoolMachine {
393+ Spec : infrav1.AzureMachinePoolMachineSpec {
394+ ProviderID : FakeProviderID ,
395+ },
396+ })
397+ params .AzureMachinePoolMachine = ampm
398+ s , err := NewMachinePoolMachineScope (params )
399+ g .Expect (err ).ToNot (HaveOccurred ())
400+ g .Expect (s ).ToNot (BeNil ())
401+ s .Logger = klogr .New ()
402+ s .workloadNodeGetter = mockClient
403+
404+ err = s .CordonAndDrain (context .TODO ())
405+ if c .Err == "" {
406+ g .Expect (err ).To (Succeed ())
407+ } else {
408+ g .Expect (err ).To (MatchError (c .Err ))
409+ }
410+ })
411+ }
412+ }
413+
319414func getReadyNode () * corev1.Node {
320415 return & corev1.Node {
321416 ObjectMeta : metav1.ObjectMeta {
0 commit comments