Issue 165: Add subnet_supplier to ec2_repository

pkg/remote/aws/repository/ec2_repository.go

@@ -17,6 +17,7 @@ type EC2Repository interface {
 	ListAllInstances() ([]*ec2.Instance, error)
 	ListAllKeyPairs() ([]*ec2.KeyPairInfo, error)
 	ListAllInternetGateways() ([]*ec2.InternetGateway, error)
+	ListAllSubnets() ([]*ec2.Subnet, []*ec2.Subnet, error)
 }
 
 type EC2Client interface {
@@ -174,3 +175,24 @@ func (r *ec2Repository) ListAllInternetGateways() ([]*ec2.InternetGateway, error
 	}
 	return internetGateways, nil
 }
+
+func (r *ec2Repository) ListAllSubnets() ([]*ec2.Subnet, []*ec2.Subnet, error) {
+	input := ec2.DescribeSubnetsInput{}
+	var subnets []*ec2.Subnet
+	var defaultSubnets []*ec2.Subnet
+	err := r.client.DescribeSubnetsPages(&input,
+		func(resp *ec2.DescribeSubnetsOutput, lastPage bool) bool {
+			for _, subnet := range resp.Subnets {
+				if subnet.DefaultForAz != nil && *subnet.DefaultForAz {
+					defaultSubnets = append(defaultSubnets, subnet)
+					continue
+				}
+				subnets = append(subnets, subnet)
+			}
+			return !lastPage
+		})
+	if err != nil {
+		return nil, nil, err
+	}
+	return subnets, defaultSubnets, nil
+}

pkg/remote/aws/repository/ec2_repository_test.go

@@ -584,3 +584,111 @@ func Test_ec2Repository_ListAllInternetGateways(t *testing.T) {
 		})
 	}
 }
+
+func Test_ec2Repository_ListAllSubnets(t *testing.T) {
+	tests := []struct {
+		name              string
+		mocks             func(client *MockEC2Client)
+		wantSubnet        []*ec2.Subnet
+		wantDefaultSubnet []*ec2.Subnet
+		wantErr           error
+	}{
+		{
+			name: "List with 2 pages",
+			mocks: func(client *MockEC2Client) {
+				client.On("DescribeSubnetsPages",
+					&ec2.DescribeSubnetsInput{},
+					mock.MatchedBy(func(callback func(res *ec2.DescribeSubnetsOutput, lastPage bool) bool) bool {
+						callback(&ec2.DescribeSubnetsOutput{
+							Subnets: []*ec2.Subnet{
+								{
+									SubnetId:     aws.String("subnet-0b13f1e0eacf67424"), // subnet2
+									DefaultForAz: aws.Bool(false),
+								},
+								{
+									SubnetId:     aws.String("subnet-0c9b78001fe186e22"), // subnet3
+									DefaultForAz: aws.Bool(false),
+								},
+								{
+									SubnetId:     aws.String("subnet-05810d3f933925f6d"), // subnet1
+									DefaultForAz: aws.Bool(false),
+								},
+							},
+						}, false)
+						callback(&ec2.DescribeSubnetsOutput{
+							Subnets: []*ec2.Subnet{
+								{
+									SubnetId:     aws.String("subnet-44fe0c65"), // us-east-1a
+									DefaultForAz: aws.Bool(true),
+								},
+								{
+									SubnetId:     aws.String("subnet-65e16628"), // us-east-1b
+									DefaultForAz: aws.Bool(true),
+								},
+								{
+									SubnetId:     aws.String("subnet-afa656f0"), // us-east-1c
+									DefaultForAz: aws.Bool(true),
+								},
+							},
+						}, true)
+						return true
+					})).Return(nil)
+			},
+			wantSubnet: []*ec2.Subnet{
+				{
+					SubnetId:     aws.String("subnet-0b13f1e0eacf67424"), // subnet2
+					DefaultForAz: aws.Bool(false),
+				},
+				{
+					SubnetId:     aws.String("subnet-0c9b78001fe186e22"), // subnet3
+					DefaultForAz: aws.Bool(false),
+				},
+				{
+					SubnetId:     aws.String("subnet-05810d3f933925f6d"), // subnet1
+					DefaultForAz: aws.Bool(false),
+				},
+			},
+			wantDefaultSubnet: []*ec2.Subnet{
+				{
+					SubnetId:     aws.String("subnet-44fe0c65"), // us-east-1a
+					DefaultForAz: aws.Bool(true),
+				},
+				{
+					SubnetId:     aws.String("subnet-65e16628"), // us-east-1b
+					DefaultForAz: aws.Bool(true),
+				},
+				{
+					SubnetId:     aws.String("subnet-afa656f0"), // us-east-1c
+					DefaultForAz: aws.Bool(true),
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			client := &MockEC2Client{}
+			tt.mocks(client)
+			r := &ec2Repository{
+				client: client,
+			}
+			gotSubnet, gotDefaultSubnet, err := r.ListAllSubnets()
+			assert.Equal(t, tt.wantErr, err)
+			changelog, err := diff.Diff(gotSubnet, tt.wantSubnet)
+			assert.Nil(t, err)
+			if len(changelog) > 0 {
+				for _, change := range changelog {
+					t.Errorf("%s: %s -> %s", strings.Join(change.Path, "."), change.From, change.To)
+				}
+				t.Fail()
+			}
+			changelog, err = diff.Diff(gotDefaultSubnet, tt.wantDefaultSubnet)
+			assert.Nil(t, err)
+			if len(changelog) > 0 {
+				for _, change := range changelog {
+					t.Errorf("%s: %s -> %s", strings.Join(change.Path, "."), change.From, change.To)
+				}
+				t.Fail()
+			}
+		})
+	}
+}

pkg/remote/aws/repository/mock_EC2Repository.go

@@ -173,6 +173,38 @@ func (_m *MockEC2Repository) ListAllSnapshots() ([]*ec2.Snapshot, error) {
 	return r0, r1
 }
 
+// ListAllSubnets provides a mock function with given fields:
+func (_m *MockEC2Repository) ListAllSubnets() ([]*ec2.Subnet, []*ec2.Subnet, error) {
+	ret := _m.Called()
+
+	var r0 []*ec2.Subnet
+	if rf, ok := ret.Get(0).(func() []*ec2.Subnet); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*ec2.Subnet)
+		}
+	}
+
+	var r1 []*ec2.Subnet
+	if rf, ok := ret.Get(1).(func() []*ec2.Subnet); ok {
+		r1 = rf()
+	} else {
+		if ret.Get(1) != nil {
+			r1 = ret.Get(1).([]*ec2.Subnet)
+		}
+	}
+
+	var r2 error
+	if rf, ok := ret.Get(2).(func() error); ok {
+		r2 = rf()
+	} else {
+		r2 = ret.Error(2)
+	}
+
+	return r0, r1, r2
+}
+
 // ListAllVolumes provides a mock function with given fields:
 func (_m *MockEC2Repository) ListAllVolumes() ([]*ec2.Volume, error) {
 	ret := _m.Called()

pkg/remote/aws/subnet_supplier.go

@@ -2,7 +2,7 @@ package aws
 
 import (
 	"github.com/aws/aws-sdk-go/service/ec2"
-	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
+	"github.com/cloudskiff/driftctl/pkg/remote/aws/repository"
 	remoteerror "github.com/cloudskiff/driftctl/pkg/remote/error"
 
 	"github.com/cloudskiff/driftctl/pkg/resource/aws"
@@ -18,7 +18,7 @@ import (
 type SubnetSupplier struct {
 	reader              terraform.ResourceReader
 	deserializer        *resource.Deserializer
-	client              ec2iface.EC2API
+	client              repository.EC2Repository
 	defaultSubnetRunner *terraform.ParallelResourceReader
 	subnetRunner        *terraform.ParallelResourceReader
 }
@@ -27,28 +27,14 @@ func NewSubnetSupplier(provider *AWSTerraformProvider, deserializer *resource.De
 	return &SubnetSupplier{
 		provider,
 		deserializer,
-		ec2.New(provider.session),
+		repository.NewEC2Repository(provider.session),
 		terraform.NewParallelResourceReader(provider.Runner().SubRunner()),
 		terraform.NewParallelResourceReader(provider.Runner().SubRunner()),
 	}
 }
 
 func (s *SubnetSupplier) Resources() ([]resource.Resource, error) {
-	input := ec2.DescribeSubnetsInput{}
+	subnets, defaultSubnets, err := s.client.ListAllSubnets()
-	var subnets []*ec2.Subnet
-	var defaultSubnets []*ec2.Subnet
-	err := s.client.DescribeSubnetsPages(&input,
-		func(resp *ec2.DescribeSubnetsOutput, lastPage bool) bool {
-			for _, subnet := range resp.Subnets {
-				if subnet.DefaultForAz != nil && *subnet.DefaultForAz {
-					defaultSubnets = append(defaultSubnets, subnet)
-					continue
-				}
-				subnets = append(subnets, subnet)
-			}
-			return !lastPage
-		},
-	)
 
 	if err != nil {
 		return nil, remoteerror.NewResourceEnumerationError(err, aws.AwsSubnetResourceType)

pkg/remote/aws/subnet_supplier_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/cloudskiff/driftctl/pkg/remote/aws/repository"
 	remoteerror "github.com/cloudskiff/driftctl/pkg/remote/error"
 	awstest "github.com/cloudskiff/driftctl/test/aws"
 	testresource "github.com/cloudskiff/driftctl/test/resource"
@@ -32,31 +33,35 @@ func TestSubnetSupplier_Resources(t *testing.T) {
 	cases := []struct {
 		test    string
 		dirName string
-		mocks   func(client *awstest.MockFakeEC2)
+		mocks   func(repo *repository.MockEC2Repository)
 		err     error
 	}{
 		{
 			test:    "no Subnet",
 			dirName: "subnet_empty",
-			mocks: func(client *awstest.MockFakeEC2) {
+			mocks: func(repo *repository.MockEC2Repository) {
-				client.On("DescribeSubnetsPages",
+				repo.On("ListAllSubnets").Once().Return([]*ec2.Subnet{}, []*ec2.Subnet{}, nil)
-					&ec2.DescribeSubnetsInput{},
-					mock.MatchedBy(func(callback func(res *ec2.DescribeSubnetsOutput, lastPage bool) bool) bool {
-						callback(&ec2.DescribeSubnetsOutput{}, true)
-						return true
-					})).Return(nil)
 			},
 			err: nil,
 		},
 		{
 			test:    "mixed default Subnet and Subnet",
 			dirName: "subnet",
-			mocks: func(client *awstest.MockFakeEC2) {
+			mocks: func(repo *repository.MockEC2Repository) {
-				client.On("DescribeSubnetsPages",
+				repo.On("ListAllSubnets").Once().Return([]*ec2.Subnet{
-					&ec2.DescribeSubnetsInput{},
+					{
-					mock.MatchedBy(func(callback func(res *ec2.DescribeSubnetsOutput, lastPage bool) bool) bool {
+						SubnetId:     aws.String("subnet-05810d3f933925f6d"), // subnet1
-						callback(&ec2.DescribeSubnetsOutput{
+						DefaultForAz: aws.Bool(false),
-							Subnets: []*ec2.Subnet{
+					},
+					{
+						SubnetId:     aws.String("subnet-0b13f1e0eacf67424"), // subnet2
+						DefaultForAz: aws.Bool(false),
+					},
+					{
+						SubnetId:     aws.String("subnet-0c9b78001fe186e22"), // subnet3
+						DefaultForAz: aws.Bool(false),
+					},
+				}, []*ec2.Subnet{
 					{
 						SubnetId:     aws.String("subnet-44fe0c65"), // us-east-1a
 						DefaultForAz: aws.Bool(true),
@@ -69,38 +74,15 @@ func TestSubnetSupplier_Resources(t *testing.T) {
 						SubnetId:     aws.String("subnet-afa656f0"), // us-east-1c
 						DefaultForAz: aws.Bool(true),
 					},
-								{
+				}, nil)
-									SubnetId:     aws.String("subnet-05810d3f933925f6d"), // subnet1
-									DefaultForAz: aws.Bool(false),
-								},
-							},
-						}, false)
-						callback(&ec2.DescribeSubnetsOutput{
-							Subnets: []*ec2.Subnet{
-								{
-									SubnetId:     aws.String("subnet-0b13f1e0eacf67424"), // subnet2
-									DefaultForAz: aws.Bool(false),
-								},
-								{
-									SubnetId:     aws.String("subnet-0c9b78001fe186e22"), // subnet3
-									DefaultForAz: aws.Bool(false),
-								},
-							},
-						}, true)
-						return true
-					})).Return(nil)
 			},
 			err: nil,
 		},
 		{
 			test:    "cannot list Subnet",
 			dirName: "subnet_empty",
-			mocks: func(client *awstest.MockFakeEC2) {
+			mocks: func(repo *repository.MockEC2Repository) {
-				client.On("DescribeSubnetsPages",
+				repo.On("ListAllSubnets").Once().Return(nil, nil, awserr.NewRequestFailure(nil, 403, ""))
-					&ec2.DescribeSubnetsInput{},
-					mock.MatchedBy(func(callback func(res *ec2.DescribeSubnetsOutput, lastPage bool) bool) bool {
-						return true
-					})).Return(awserr.NewRequestFailure(nil, 403, ""))
 			},
 			err: remoteerror.NewResourceEnumerationError(awserr.NewRequestFailure(nil, 403, ""), resourceaws.AwsSubnetResourceType),
 		},
@@ -125,7 +107,7 @@ func TestSubnetSupplier_Resources(t *testing.T) {
 		}
 
 		t.Run(c.test, func(tt *testing.T) {
-			fakeEC2 := awstest.MockFakeEC2{}
+			fakeEC2 := repository.MockEC2Repository{}
 			c.mocks(&fakeEC2)
 			provider := mocks2.NewMockedGoldenTFProvider(c.dirName, providerLibrary.Provider(terraform.AWS), shouldUpdate)
 			s := &SubnetSupplier{