From f625b4a2d9be9b48ef11ff27db497f8a80fcd181 Mon Sep 17 00:00:00 2001
From: Ashley Gau <agau4779@gmail.com>
Date: Mon, 26 Feb 2018 14:18:11 -0800
Subject: [PATCH] add hooks to add, remove, insert instances from
 instancegroups

---
 .../providers/gce/cloud/mock/mock.go          | 107 +++++++++++++++++-
 .../gce/gce_loadbalancer_test_utils.go        |   4 +
 2 files changed, 110 insertions(+), 1 deletion(-)

diff --git a/pkg/cloudprovider/providers/gce/cloud/mock/mock.go b/pkg/cloudprovider/providers/gce/cloud/mock/mock.go
index 06ff480d47b..8d9b7382efd 100644
--- a/pkg/cloudprovider/providers/gce/cloud/mock/mock.go
+++ b/pkg/cloudprovider/providers/gce/cloud/mock/mock.go
@@ -33,7 +33,8 @@ import (
 	beta "google.golang.org/api/compute/v0.beta"
 	ga "google.golang.org/api/compute/v1"
 	"google.golang.org/api/googleapi"
-	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud"
+	cloud "k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud"
+	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud/filter"
 	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud/meta"
 )
 
@@ -226,3 +227,107 @@ func InsertAlphaAddressHook(ctx context.Context, key *meta.Key, obj *alpha.Addre
 	projectID := m.ProjectRouter.ProjectID(ctx, meta.VersionBeta, "addresses")
 	return convertAndInsertAlphaAddress(key, obj, m.Objects, meta.VersionAlpha, projectID)
 }
+
+// Map from InstanceGroup key to list of Instances
+type instanceGroupAttributes struct {
+	instanceMap map[meta.Key][]*ga.InstanceWithNamedPorts
+}
+
+func AddInstancesHook(ctx context.Context, key *meta.Key, req *ga.InstanceGroupsAddInstancesRequest, m *cloud.MockInstanceGroups) error {
+	_, err := m.Get(ctx, key)
+	if err != nil {
+		return &googleapi.Error{
+			Code:    http.StatusNotFound,
+			Message: fmt.Sprintf("Key: %s was not found in InstanceGroups", key.String()),
+		}
+	}
+
+	if m.X == nil {
+		m.X = instanceGroupAttributes{
+			instanceMap: make(map[meta.Key][]*ga.InstanceWithNamedPorts),
+		}
+	}
+
+	var attrs instanceGroupAttributes
+	attrs = m.X.(instanceGroupAttributes)
+
+	instances, ok := attrs.instanceMap[*key]
+	if !ok {
+		instances = []*ga.InstanceWithNamedPorts{}
+	}
+
+	for _, instance := range req.Instances {
+		iWithPort := &ga.InstanceWithNamedPorts{
+			Instance: instance.Instance,
+		}
+
+		instances = append(instances, iWithPort)
+	}
+
+	attrs.instanceMap[*key] = instances
+	m.X = attrs
+	return nil
+}
+
+func ListInstancesHook(ctx context.Context, key *meta.Key, req *ga.InstanceGroupsListInstancesRequest, filter *filter.F, m *cloud.MockInstanceGroups) ([]*ga.InstanceWithNamedPorts, error) {
+	_, err := m.Get(ctx, key)
+	if err != nil {
+		return nil, &googleapi.Error{
+			Code:    http.StatusNotFound,
+			Message: fmt.Sprintf("Key: %s was not found in InstanceGroups", key.String()),
+		}
+	}
+
+	if m.X == nil {
+		m.X = instanceGroupAttributes{instanceMap: make(map[meta.Key][]*ga.InstanceWithNamedPorts)}
+	}
+
+	var attrs instanceGroupAttributes
+	attrs = m.X.(instanceGroupAttributes)
+
+	instances, ok := attrs.instanceMap[*key]
+	if !ok {
+		instances = []*ga.InstanceWithNamedPorts{}
+	}
+
+	return instances, nil
+}
+
+func RemoveInstancesHook(ctx context.Context, key *meta.Key, req *ga.InstanceGroupsRemoveInstancesRequest, m *cloud.MockInstanceGroups) error {
+	_, err := m.Get(ctx, key)
+	if err != nil {
+		return &googleapi.Error{
+			Code:    http.StatusNotFound,
+			Message: fmt.Sprintf("Key: %s was not found in InstanceGroups", key.String()),
+		}
+	}
+
+	if m.X == nil {
+		m.X = instanceGroupAttributes{
+			instanceMap: make(map[meta.Key][]*ga.InstanceWithNamedPorts),
+		}
+	}
+
+	var attrs instanceGroupAttributes
+	attrs = m.X.(instanceGroupAttributes)
+
+	instances, ok := attrs.instanceMap[*key]
+	if !ok {
+		instances = []*ga.InstanceWithNamedPorts{}
+	}
+
+	for _, instanceToRemove := range req.Instances {
+		for i, instance := range instances {
+			if instanceToRemove.Instance == instance.Instance {
+				// Delete instance from instances without preserving order
+				instances[i] = instances[len(instances)-1]
+				instances = instances[:len(instances)-1]
+				break
+			}
+		}
+	}
+
+	attrs.instanceMap[*key] = instances
+	m.X = attrs
+	return nil
+}
diff --git a/pkg/cloudprovider/providers/gce/gce_loadbalancer_test_utils.go b/pkg/cloudprovider/providers/gce/gce_loadbalancer_test_utils.go
index 8e57e35ad12..67ee161350f 100644
--- a/pkg/cloudprovider/providers/gce/gce_loadbalancer_test_utils.go
+++ b/pkg/cloudprovider/providers/gce/gce_loadbalancer_test_utils.go
@@ -84,6 +84,10 @@ func fakeGCECloud(projectID, region, zoneName string) (*GCECloud, error) {
 	c.MockAddresses.InsertHook = mock.InsertAddressHook
 	c.MockAlphaAddresses.InsertHook = mock.InsertAlphaAddressHook
 
+	c.MockInstanceGroups.AddInstancesHook = mock.AddInstancesHook
+	c.MockInstanceGroups.RemoveInstancesHook = mock.RemoveInstancesHook
+	c.MockInstanceGroups.ListInstancesHook = mock.ListInstancesHook
+
 	keyGA := meta.GlobalKey("key-ga")
 	c.MockZones.Objects[*keyGA] = &cloud.MockZonesObj{
 		Obj: &compute.Zone{Name: zoneName, Region: gce.getRegionLink(region)},