Merge pull request #5072 from justinsb/fix_aws_test

pkg/cloudprovider/aws/aws.go

@@ -33,8 +33,10 @@ import (
 	"github.com/golang/glog"
 )
 
+// Abstraction over EC2, to allow mocking/other implementations
 type EC2 interface {
-	Instances(instIds []string, filter *ec2.Filter) (resp *ec2.InstancesResp, err error)
+	// Query EC2 for instances matching the filter
+	Instances(instIds []string, filter *ec2InstanceFilter) (resp *ec2.InstancesResp, err error)
 }
 
 // AWSCloud is an implementation of Interface, TCPLoadBalancer and Instances for Amazon Web Services.
@@ -49,6 +51,37 @@ type AWSCloudConfig struct {
 	}
 }
 
+// Similar to ec2.Filter, but the filter values can be read from tests
+// (ec2.Filter only has private members)
+type ec2InstanceFilter struct {
+	PrivateDNSName string
+}
+
+// True if the passed instance matches the filter
+func (f *ec2InstanceFilter) Matches(instance ec2.Instance) bool {
+	if f.PrivateDNSName != "" && instance.PrivateDNSName != f.PrivateDNSName {
+		return false
+	}
+	return true
+}
+
+// goamzEC2 is an implementation of the EC2 interface, backed by goamz
+type GoamzEC2 struct {
+	ec2 *ec2.EC2
+}
+
+// Implementation of EC2.Instances
+func (self *GoamzEC2) Instances(instanceIds []string, filter *ec2InstanceFilter) (resp *ec2.InstancesResp, err error) {
+	var goamzFilter *ec2.Filter
+	if filter != nil {
+		goamzFilter = ec2.NewFilter()
+		if filter.PrivateDNSName != "" {
+			goamzFilter.Add("private-dns-name", filter.PrivateDNSName)
+		}
+	}
+	return self.ec2.Instances(instanceIds, goamzFilter)
+}
+
 type AuthFunc func() (auth aws.Auth, err error)
 
 func init() {
@@ -99,7 +132,7 @@ func newAWSCloud(config io.Reader, authFunc AuthFunc) (*AWSCloud, error) {
 
 	ec2 := ec2.New(auth, region)
 	return &AWSCloud{
-		ec2: ec2,
+		ec2: &GoamzEC2{ec2: ec2},
 		cfg: cfg,
 	}, nil
 }
@@ -147,8 +180,8 @@ func (aws *AWSCloud) ExternalID(name string) (string, error) {
 
 // Return the instances matching the relevant private dns name.
 func (aws *AWSCloud) getInstancesByDnsName(name string) (*ec2.Instance, error) {
-	f := ec2.NewFilter()
-	f.Add("private-dns-name", name)
+	f := &ec2InstanceFilter{}
+	f.PrivateDNSName = name
 
 	resp, err := aws.ec2.Instances(nil, f)
 	if err != nil {

pkg/cloudprovider/aws/aws_test.go

@@ -76,20 +76,26 @@ func TestNewAWSCloud(t *testing.T) {
 }
 
 type FakeEC2 struct {
-	instances func(instanceIds []string, filter *ec2.Filter) (resp *ec2.InstancesResp, err error)
+	instances func(instanceIds []string, filter *ec2InstanceFilter) (resp *ec2.InstancesResp, err error)
 }
 
-func (ec2 *FakeEC2) Instances(instanceIds []string, filter *ec2.Filter) (resp *ec2.InstancesResp, err error) {
+func (ec2 *FakeEC2) Instances(instanceIds []string, filter *ec2InstanceFilter) (resp *ec2.InstancesResp, err error) {
 	return ec2.instances(instanceIds, filter)
 }
 
 func mockInstancesResp(instances []ec2.Instance) (aws *AWSCloud) {
 	return &AWSCloud{
 		&FakeEC2{
-			func(instanceIds []string, filter *ec2.Filter) (resp *ec2.InstancesResp, err error) {
+			func(instanceIds []string, filter *ec2InstanceFilter) (resp *ec2.InstancesResp, err error) {
+				matches := []ec2.Instance{}
+				for _, instance := range instances {
+					if filter == nil || filter.Matches(instance) {
+						matches = append(matches, instance)
+					}
+				}
 				return &ec2.InstancesResp{"",
 					[]ec2.Reservation{
-						{"", "", "", nil, instances}}}, nil
+						{"", "", "", nil, matches}}}, nil
 			}},
 		nil}
 }
@@ -128,10 +134,12 @@ func TestList(t *testing.T) {
 }
 
 func TestIPAddress(t *testing.T) {
+	// Note these instances have the same name
+	// (we test that this produces an error)
 	instances := make([]ec2.Instance, 2)
 	instances[0].PrivateDNSName = "instance1"
 	instances[0].PrivateIpAddress = "192.168.0.1"
-	instances[1].PrivateDNSName = "instance2"
+	instances[1].PrivateDNSName = "instance1"
 	instances[1].PrivateIpAddress = "192.168.0.2"
 
 	aws1 := mockInstancesResp([]ec2.Instance{})