From f92f2f4782034398fa1f38258151bba0000118ba Mon Sep 17 00:00:00 2001
From: Vikas Choudhary <vichoudh@redhat.com>
Date: Fri, 17 May 2019 20:30:14 +0530
Subject: [PATCH] Update machine.Spec.ProviderID field

---
 pkg/actuators/machine/actuator.go      | 33 +++++++++++++++++++++++++-
 pkg/actuators/machine/actuator_test.go |  5 +++-
 pkg/actuators/machine/stubs.go         |  4 ++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/pkg/actuators/machine/actuator.go b/pkg/actuators/machine/actuator.go
index 953878e876..305248e8b0 100644
--- a/pkg/actuators/machine/actuator.go
+++ b/pkg/actuators/machine/actuator.go
@@ -37,6 +37,7 @@ import (
 	apierrors "github.com/openshift/cluster-api/pkg/errors"
 	providerconfigv1 "sigs.k8s.io/cluster-api-provider-aws/pkg/apis/awsproviderconfig/v1beta1"
 
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 
 	awsclient "sigs.k8s.io/cluster-api-provider-aws/pkg/client"
@@ -116,7 +117,37 @@ func (a *Actuator) Create(context context.Context, cluster *clusterv1.Cluster, m
 		}
 		return err
 	}
-	return a.updateStatus(machine, instance)
+	updatedMachine, err := a.updateProviderID(machine, instance)
+	if err != nil {
+		return fmt.Errorf("failed to update machine object with providerID: %v", err)
+	}
+	return a.updateStatus(updatedMachine, instance)
+}
+
+// updateProviderID adds providerID in the machine spec
+func (a *Actuator) updateProviderID(machine *machinev1.Machine, instance *ec2.Instance) (*machinev1.Machine, error) {
+	existingProviderID := machine.Spec.ProviderID
+	machineCopy := machine.DeepCopy()
+	if instance != nil {
+		providerID := fmt.Sprintf("aws:///%s/%s", aws.StringValue(instance.Placement.AvailabilityZone), aws.StringValue(instance.InstanceId))
+
+		if existingProviderID != nil && *existingProviderID == providerID {
+			glog.Infof("ProviderID already set in the machine Spec with value:%s", *existingProviderID)
+			return machine, nil
+		}
+		machineCopy.Spec.ProviderID = &providerID
+		if err := a.client.Update(context.Background(), machineCopy); err != nil {
+			return nil, fmt.Errorf("error updating machine spec ProviderID: %v", err)
+		}
+		glog.Infof("ProviderID updated at machine spec: %s", providerID)
+	} else {
+		machineCopy.Spec.ProviderID = nil
+		if err := a.client.Update(context.Background(), machineCopy); err != nil {
+			return nil, fmt.Errorf("error updating ProviderID in machine spec: %v", err)
+		}
+		glog.Infof("No instance found so clearing ProviderID field in the machine spec")
+	}
+	return machineCopy, nil
 }
 
 func (a *Actuator) updateMachineStatus(machine *machinev1.Machine, awsStatus *providerconfigv1.AWSMachineProviderStatus, networkAddresses []corev1.NodeAddress) error {
diff --git a/pkg/actuators/machine/actuator_test.go b/pkg/actuators/machine/actuator_test.go
index 3bcb95d719..bbe5a46042 100644
--- a/pkg/actuators/machine/actuator_test.go
+++ b/pkg/actuators/machine/actuator_test.go
@@ -642,7 +642,7 @@ func TestAvailabiltyZone(t *testing.T) {
 			if tc.availabilityZone != "" && tc.subnet == "" {
 				placement = &ec2.Placement{AvailabilityZone: aws.String(tc.availabilityZone)}
 			}
-
+			az := "us-east-1a"
 			mockAWSClient.EXPECT().RunInstances(placementMatcher{placement}).Return(
 				&ec2.Reservation{
 					Instances: []*ec2.Instance{
@@ -653,6 +653,9 @@ func TestAvailabiltyZone(t *testing.T) {
 								Name: aws.String("Running"),
 							},
 							LaunchTime: aws.Time(time.Now()),
+							Placement: &ec2.Placement{
+								AvailabilityZone: &az,
+							},
 						},
 					},
 				}, nil)
diff --git a/pkg/actuators/machine/stubs.go b/pkg/actuators/machine/stubs.go
index be56557031..73080d5fc5 100644
--- a/pkg/actuators/machine/stubs.go
+++ b/pkg/actuators/machine/stubs.go
@@ -241,6 +241,7 @@ func stubDescribeTargetGroupsOutput() *elbv2.DescribeTargetGroupsOutput {
 }
 
 func stubReservation(imageID, instanceID string) *ec2.Reservation {
+	az := defaultAvailabilityZone
 	return &ec2.Reservation{
 		Instances: []*ec2.Instance{
 			{
@@ -251,6 +252,9 @@ func stubReservation(imageID, instanceID string) *ec2.Reservation {
 					Code: aws.Int64(16),
 				},
 				LaunchTime: aws.Time(time.Now()),
+				Placement: &ec2.Placement{
+					AvailabilityZone: &az,
+				},
 			},
 		},
 	}