Refactor mocking to split functionality across packages

pkg/cloudprovider/aws/builder.go

@@ -4,13 +4,12 @@ import (
 	"fmt"
 
 	"github.com/atlassian-labs/cyclops/pkg/cloudprovider"
+	fakeaws "github.com/atlassian-labs/cyclops/pkg/cloudprovider/aws/fake"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/autoscaling"
-	"github.com/aws/aws-sdk-go/service/autoscaling/autoscalingiface"
 	"github.com/aws/aws-sdk-go/service/ec2"
-	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/go-logr/logr"
 )
 
@@ -43,8 +42,16 @@ func NewCloudProvider(logger logr.Logger) (cloudprovider.CloudProvider, error) {
 	return p, nil
 }
 
-// NewMockCloudProvider returns a new mock AWS cloud provider
-func NewMockCloudProvider(autoscalingiface autoscalingiface.AutoScalingAPI, ec2iface ec2iface.EC2API) cloudprovider.CloudProvider {
+// NewFakeCloudProvider returns a new mock AWS cloud provider
+func NewFakeCloudProvider(instances *[]*fakeaws.Instance) cloudprovider.CloudProvider {
+	autoscalingiface := &fakeaws.Autoscaling{
+		Instances: instances,
+	}
+
+	ec2iface := &fakeaws.Ec2{
+		Instances: instances,
+	}
+
 	return &provider{
 		autoScalingService: autoscalingiface,
 		ec2Service:         ec2iface,

pkg/cloudprovider/aws/fake/fake.go

@@ -0,0 +1,142 @@
+package fakeaws
+
+import (
+	"fmt"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/autoscaling"
+	"github.com/aws/aws-sdk-go/service/ec2"
+
+	"github.com/aws/aws-sdk-go/service/autoscaling/autoscalingiface"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
+)
+
+var (
+	defaultAvailabilityZone = "us-east-1a"
+)
+
+type Instance struct {
+	InstanceID           string
+	AutoscalingGroupName string
+	State                string
+}
+
+type Ec2 struct {
+	ec2iface.EC2API
+
+	Instances *[]*Instance
+}
+
+type Autoscaling struct {
+	autoscalingiface.AutoScalingAPI
+
+	Instances *[]*Instance
+}
+
+func GenerateProviderID(instanceID string) string {
+	return fmt.Sprintf("aws:///%s/%s",
+		defaultAvailabilityZone,
+		instanceID,
+	)
+}
+
+func generateEc2Instance(instance *Instance) *ec2.Instance {
+	ec2Instance := &ec2.Instance{
+		InstanceId: aws.String(instance.InstanceID),
+		State: &ec2.InstanceState{
+			Name: aws.String(instance.State),
+		},
+	}
+
+	return ec2Instance
+}
+
+func generateAutoscalingInstance(instance *Instance) *autoscaling.Instance {
+	autoscalingInstance := &autoscaling.Instance{
+		InstanceId:       aws.String(instance.InstanceID),
+		AvailabilityZone: aws.String(defaultAvailabilityZone),
+	}
+
+	return autoscalingInstance
+}
+
+// *************** Autoscaling *************** //
+
+func (m *Autoscaling) DescribeAutoScalingGroups(input *autoscaling.DescribeAutoScalingGroupsInput) (*autoscaling.DescribeAutoScalingGroupsOutput, error) {
+	var asgs = make(map[string]*autoscaling.Group, 0)
+
+	var asgNameLookup = make(map[string]interface{})
+
+	for _, asgName := range input.AutoScalingGroupNames {
+		asgNameLookup[*asgName] = nil
+	}
+
+	for _, instance := range *m.Instances {
+		if instance.AutoscalingGroupName == "" {
+			continue
+		}
+
+		if _, exists := asgNameLookup[instance.AutoscalingGroupName]; !exists {
+			continue
+		}
+
+		asg, exists := asgs[instance.AutoscalingGroupName]
+
+		if !exists {
+			asg = &autoscaling.Group{
+				AutoScalingGroupName: aws.String(instance.AutoscalingGroupName),
+				Instances:            []*autoscaling.Instance{},
+				AvailabilityZones: []*string{
+					aws.String(defaultAvailabilityZone),
+				},
+			}
+
+			asgs[instance.AutoscalingGroupName] = asg
+		}
+
+		asg.Instances = append(
+			asg.Instances,
+			generateAutoscalingInstance(instance),
+		)
+	}
+
+	var asgList = make([]*autoscaling.Group, 0)
+
+	for _, asg := range asgs {
+		asgList = append(asgList, asg)
+	}
+
+	return &autoscaling.DescribeAutoScalingGroupsOutput{
+		AutoScalingGroups: asgList,
+	}, nil
+}
+
+// *************** EC2 *************** //
+
+func (m *Ec2) DescribeInstances(input *ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error) {
+	var instances = make([]*ec2.Instance, 0)
+	var instanceIds = make(map[string]interface{})
+
+	for _, instanceId := range input.InstanceIds {
+		instanceIds[*instanceId] = nil
+	}
+
+	for _, instance := range *m.Instances {
+		if _, ok := instanceIds[instance.InstanceID]; input.InstanceIds != nil && !ok {
+			continue
+		}
+
+		instances = append(
+			instances,
+			generateEc2Instance(instance),
+		)
+	}
+
+	return &ec2.DescribeInstancesOutput{
+		Reservations: []*ec2.Reservation{
+			{
+				Instances: instances,
+			},
+		},
+	}, nil
+}

pkg/controller/cyclenoderequest/transitioner/fake/fake.go

@@ -0,0 +1,46 @@
+package faketransitioner
+
+import (
+	"net/http"
+
+	v1 "github.com/atlassian-labs/cyclops/pkg/apis/atlassian/v1"
+	"github.com/atlassian-labs/cyclops/pkg/controller"
+	"github.com/atlassian-labs/cyclops/pkg/controller/cyclenoderequest/transitioner"
+	"github.com/atlassian-labs/cyclops/pkg/mock"
+)
+
+type Transitioner struct {
+	*transitioner.CycleNodeRequestTransitioner
+	*mock.Client
+
+	cloudProviderInstances []*mock.Node
+	kubeNodes              []*mock.Node
+}
+
+func NewFakeTransitioner(cnr *v1.CycleNodeRequest, opts ...Option) *Transitioner {
+	t := &Transitioner{
+		// By default there are no nodes and each test will
+		// override these as needed
+		cloudProviderInstances: make([]*mock.Node, 0),
+		kubeNodes:              make([]*mock.Node, 0),
+	}
+
+	for _, opt := range opts {
+		opt(t)
+	}
+
+	t.Client = mock.NewClient(t.kubeNodes, t.cloudProviderInstances, cnr)
+
+	rm := &controller.ResourceManager{
+		Client:        t.K8sClient,
+		RawClient:     t.RawClient,
+		HttpClient:    http.DefaultClient,
+		CloudProvider: t.CloudProvider,
+	}
+
+	t.CycleNodeRequestTransitioner = transitioner.NewCycleNodeRequestTransitioner(
+		cnr, rm, transitioner.Options{},
+	)
+
+	return t
+}

pkg/controller/cyclenoderequest/transitioner/fake/options.go

@@ -0,0 +1,17 @@
+package faketransitioner
+
+import "github.com/atlassian-labs/cyclops/pkg/mock"
+
+type Option func(t *Transitioner)
+
+func WithCloudProviderInstances(nodes []*mock.Node) Option {
+	return func(t *Transitioner) {
+		t.cloudProviderInstances = append(t.cloudProviderInstances, nodes...)
+	}
+}
+
+func WithKubeNodes(nodes []*mock.Node) Option {
+	return func(t *Transitioner) {
+		t.kubeNodes = append(t.kubeNodes, nodes...)
+	}
+}