From cc81742e4097e43d9a55b720c1b3d1598313a145 Mon Sep 17 00:00:00 2001
From: Justin Santa Barbara <justin@fathomdb.com>
Date: Thu, 5 Mar 2015 09:05:11 -0500
Subject: [PATCH] Implement EC2 filter in tests

---
 pkg/cloudprovider/aws/aws.go      | 41 ++++++++++++++++++++++++++++---
 pkg/cloudprovider/aws/aws_test.go | 14 ++++++++---
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/pkg/cloudprovider/aws/aws.go b/pkg/cloudprovider/aws/aws.go
index 41e8207865..cde318bb77 100644
--- a/pkg/cloudprovider/aws/aws.go
+++ b/pkg/cloudprovider/aws/aws.go
@@ -30,8 +30,10 @@ import (
 	"github.com/GoogleCloudPlatform/kubernetes/pkg/cloudprovider"
 )
 
+// 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.
@@ -46,6 +48,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() {
@@ -96,7 +129,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
 }
@@ -144,8 +177,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 {
diff --git a/pkg/cloudprovider/aws/aws_test.go b/pkg/cloudprovider/aws/aws_test.go
index 778de9cdf4..6df2cfbef1 100644
--- a/pkg/cloudprovider/aws/aws_test.go
+++ b/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}
 }