Skip to content

Commit 2977d63

Browse files
wbeuilwbeuil
authored
Merge pull request #1599 from martin-vanta/ml/aws-lb-listener-cache
fix: aws_lb_listener cache key
2 parents 9d35f9c + 4f9fed2 commit 2977d63

File tree

2 files changed

+142
-70
lines changed

2 files changed

+142
-70
lines changed

‎enumeration/remote/aws/repository/elbv2_repository.go

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package repository
22

33
import (
4+
"fmt"
5+
46
"github.com/aws/aws-sdk-go/aws/session"
57
"github.com/aws/aws-sdk-go/service/elbv2"
68
"github.com/aws/aws-sdk-go/service/elbv2/elbv2iface"
@@ -45,7 +47,8 @@ func (r *elbv2Repository) ListAllLoadBalancers() ([]*elbv2.LoadBalancer, error)
4547
}
4648

4749
func (r *elbv2Repository) ListAllLoadBalancerListeners(loadBalancerArn string) ([]*elbv2.Listener, error) {
48-
if v := r.cache.Get("elbv2ListAllLoadBalancerListeners"); v != nil {
50+
cacheKey := fmt.Sprintf("elbv2ListAllLoadBalancerListeners_%s", loadBalancerArn)
51+
if v := r.cache.Get(cacheKey); v != nil {
4952
return v.([]*elbv2.Listener), nil
5053
}
5154

@@ -60,6 +63,6 @@ func (r *elbv2Repository) ListAllLoadBalancerListeners(loadBalancerArn string) (
6063
if err != nil {
6164
return nil, err
6265
}
63-
r.cache.Put("elbv2ListAllLoadBalancerListeners", results)
66+
r.cache.Put(cacheKey, results)
6467
return results, err
6568
}

‎enumeration/remote/aws/repository/elbv2_repository_test.go

Lines changed: 137 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,17 @@
11
package repository
22

33
import (
4-
"github.com/snyk/driftctl/enumeration/remote/cache"
54
"strings"
65
"testing"
76

87
"github.com/aws/aws-sdk-go/aws"
98
"github.com/aws/aws-sdk-go/service/elbv2"
109
"github.com/pkg/errors"
11-
awstest "github.com/snyk/driftctl/test/aws"
12-
"github.com/stretchr/testify/mock"
13-
1410
"github.com/r3labs/diff/v2"
11+
"github.com/snyk/driftctl/enumeration/remote/cache"
12+
awstest "github.com/snyk/driftctl/test/aws"
1513
"github.com/stretchr/testify/assert"
14+
"github.com/stretchr/testify/mock"
1615
)
1716

1817
func Test_ELBV2Repository_ListAllLoadBalancers(t *testing.T) {
@@ -133,17 +132,53 @@ func Test_ELBV2Repository_ListAllLoadBalancers(t *testing.T) {
133132
func Test_ELBV2Repository_ListAllLoadBalancerListeners(t *testing.T) {
134133
dummyError := errors.New("dummy error")
135134

135+
type call struct {
136+
loadBalancerArn string
137+
mocks func(*awstest.MockFakeELBV2, *cache.MockCache)
138+
want []*elbv2.Listener
139+
wantErr error
140+
}
141+
136142
tests := []struct {
137-
name string
138-
mocks func(*awstest.MockFakeELBV2, *cache.MockCache)
139-
want []*elbv2.Listener
140-
wantErr error
143+
name string
144+
calls []call
141145
}{
142146
{
143147
name: "list load balancer listeners",
144-
mocks: func(client *awstest.MockFakeELBV2, store *cache.MockCache) {
145-
results := &elbv2.DescribeListenersOutput{
146-
Listeners: []*elbv2.Listener{
148+
calls: []call{
149+
{
150+
loadBalancerArn: "test-lb",
151+
mocks: func(client *awstest.MockFakeELBV2, store *cache.MockCache) {
152+
results := &elbv2.DescribeListenersOutput{
153+
Listeners: []*elbv2.Listener{
154+
{
155+
LoadBalancerArn: aws.String("test-lb"),
156+
ListenerArn: aws.String("test-lb-listener-1"),
157+
},
158+
{
159+
LoadBalancerArn: aws.String("test-lb"),
160+
ListenerArn: aws.String("test-lb-listener-2"),
161+
},
162+
},
163+
}
164+
165+
store.On("Get", "elbv2ListAllLoadBalancerListeners_test-lb").Return(nil).Once()
166+
167+
client.On("DescribeListenersPages",
168+
&elbv2.DescribeListenersInput{LoadBalancerArn: aws.String("test-lb")},
169+
mock.MatchedBy(func(callback func(res *elbv2.DescribeListenersOutput, lastPage bool) bool) bool {
170+
callback(&elbv2.DescribeListenersOutput{Listeners: []*elbv2.Listener{
171+
results.Listeners[0],
172+
}}, false)
173+
callback(&elbv2.DescribeListenersOutput{Listeners: []*elbv2.Listener{
174+
results.Listeners[1],
175+
}}, true)
176+
return true
177+
})).Return(nil).Once()
178+
179+
store.On("Put", "elbv2ListAllLoadBalancerListeners_test-lb", results.Listeners).Return(false).Once()
180+
},
181+
want: []*elbv2.Listener{
147182
{
148183
LoadBalancerArn: aws.String("test-lb"),
149184
ListenerArn: aws.String("test-lb-listener-1"),
@@ -153,90 +188,124 @@ func Test_ELBV2Repository_ListAllLoadBalancerListeners(t *testing.T) {
153188
ListenerArn: aws.String("test-lb-listener-2"),
154189
},
155190
},
156-
}
157-
158-
store.On("Get", "elbv2ListAllLoadBalancerListeners").Return(nil).Once()
159-
160-
client.On("DescribeListenersPages",
161-
&elbv2.DescribeListenersInput{LoadBalancerArn: aws.String("test-lb")},
162-
mock.MatchedBy(func(callback func(res *elbv2.DescribeListenersOutput, lastPage bool) bool) bool {
163-
callback(&elbv2.DescribeListenersOutput{Listeners: []*elbv2.Listener{
164-
results.Listeners[0],
165-
}}, false)
166-
callback(&elbv2.DescribeListenersOutput{Listeners: []*elbv2.Listener{
167-
results.Listeners[1],
168-
}}, true)
169-
return true
170-
})).Return(nil).Once()
171-
172-
store.On("Put", "elbv2ListAllLoadBalancerListeners", results.Listeners).Return(false).Once()
173-
},
174-
want: []*elbv2.Listener{
175-
{
176-
LoadBalancerArn: aws.String("test-lb"),
177-
ListenerArn: aws.String("test-lb-listener-1"),
178-
},
179-
{
180-
LoadBalancerArn: aws.String("test-lb"),
181-
ListenerArn: aws.String("test-lb-listener-2"),
182191
},
183192
},
184193
},
185194
{
186195
name: "list load balancer listeners from cache",
187-
mocks: func(client *awstest.MockFakeELBV2, store *cache.MockCache) {
188-
output := &elbv2.DescribeListenersOutput{
189-
Listeners: []*elbv2.Listener{
196+
calls: []call{
197+
{
198+
loadBalancerArn: "test-lb",
199+
mocks: func(client *awstest.MockFakeELBV2, store *cache.MockCache) {
200+
output := &elbv2.DescribeListenersOutput{
201+
Listeners: []*elbv2.Listener{
202+
{
203+
LoadBalancerArn: aws.String("test-lb"),
204+
ListenerArn: aws.String("test-lb-listener"),
205+
},
206+
},
207+
}
208+
209+
store.On("Get", "elbv2ListAllLoadBalancerListeners_test-lb").Return(output.Listeners).Once()
210+
},
211+
want: []*elbv2.Listener{
190212
{
191213
LoadBalancerArn: aws.String("test-lb"),
192214
ListenerArn: aws.String("test-lb-listener"),
193215
},
194216
},
195-
}
196-
197-
store.On("Get", "elbv2ListAllLoadBalancerListeners").Return(output.Listeners).Once()
217+
},
198218
},
199-
want: []*elbv2.Listener{
219+
},
220+
{
221+
name: "list load balancer listeners from multiple load balancers",
222+
calls: []call{
223+
{
224+
loadBalancerArn: "test-lb-1",
225+
mocks: func(client *awstest.MockFakeELBV2, store *cache.MockCache) {
226+
output := &elbv2.DescribeListenersOutput{
227+
Listeners: []*elbv2.Listener{
228+
{
229+
LoadBalancerArn: aws.String("test-lb-1"),
230+
ListenerArn: aws.String("test-lb-1-listener"),
231+
},
232+
},
233+
}
234+
235+
store.On("Get", "elbv2ListAllLoadBalancerListeners_test-lb-1").Return(output.Listeners).Once()
236+
},
237+
want: []*elbv2.Listener{
238+
{
239+
LoadBalancerArn: aws.String("test-lb-1"),
240+
ListenerArn: aws.String("test-lb-1-listener"),
241+
},
242+
},
243+
},
200244
{
201-
LoadBalancerArn: aws.String("test-lb"),
202-
ListenerArn: aws.String("test-lb-listener"),
245+
loadBalancerArn: "test-lb-2",
246+
mocks: func(client *awstest.MockFakeELBV2, store *cache.MockCache) {
247+
output := &elbv2.DescribeListenersOutput{
248+
Listeners: []*elbv2.Listener{
249+
{
250+
LoadBalancerArn: aws.String("test-lb-2"),
251+
ListenerArn: aws.String("test-lb-2-listener"),
252+
},
253+
},
254+
}
255+
256+
store.On("Get", "elbv2ListAllLoadBalancerListeners_test-lb-2").Return(output.Listeners).Once()
257+
},
258+
want: []*elbv2.Listener{
259+
{
260+
LoadBalancerArn: aws.String("test-lb-2"),
261+
ListenerArn: aws.String("test-lb-2-listener"),
262+
},
263+
},
203264
},
204265
},
205266
},
206267
{
207268
name: "error listing load balancer listeners",
208-
mocks: func(client *awstest.MockFakeELBV2, store *cache.MockCache) {
209-
store.On("Get", "elbv2ListAllLoadBalancerListeners").Return(nil).Once()
269+
calls: []call{
270+
{
271+
loadBalancerArn: "test-lb",
272+
mocks: func(client *awstest.MockFakeELBV2, store *cache.MockCache) {
273+
store.On("Get", "elbv2ListAllLoadBalancerListeners_test-lb").Return(nil).Once()
210274

211-
client.On("DescribeListenersPages",
212-
&elbv2.DescribeListenersInput{LoadBalancerArn: aws.String("test-lb")},
213-
mock.MatchedBy(func(callback func(res *elbv2.DescribeListenersOutput, lastPage bool) bool) bool {
214-
callback(&elbv2.DescribeListenersOutput{Listeners: []*elbv2.Listener{}}, true)
215-
return true
216-
})).Return(dummyError).Once()
275+
client.On("DescribeListenersPages",
276+
&elbv2.DescribeListenersInput{LoadBalancerArn: aws.String("test-lb")},
277+
mock.MatchedBy(func(callback func(res *elbv2.DescribeListenersOutput, lastPage bool) bool) bool {
278+
callback(&elbv2.DescribeListenersOutput{Listeners: []*elbv2.Listener{}}, true)
279+
return true
280+
})).Return(dummyError).Once()
281+
},
282+
wantErr: dummyError,
283+
},
217284
},
218-
wantErr: dummyError,
219285
},
220286
}
221287
for _, tt := range tests {
222288
t.Run(tt.name, func(t *testing.T) {
223289
store := &cache.MockCache{}
224290
client := &awstest.MockFakeELBV2{}
225-
tt.mocks(client, store)
226-
r := &elbv2Repository{
227-
client: client,
228-
cache: store,
229-
}
230-
got, err := r.ListAllLoadBalancerListeners("test-lb")
231-
assert.Equal(t, tt.wantErr, err)
232291

233-
changelog, err := diff.Diff(got, tt.want)
234-
assert.Nil(t, err)
235-
if len(changelog) > 0 {
236-
for _, change := range changelog {
237-
t.Errorf("%s: %v -> %v", strings.Join(change.Path, "."), change.From, change.To)
292+
for _, call := range tt.calls {
293+
call.mocks(client, store)
294+
r := &elbv2Repository{
295+
client: client,
296+
cache: store,
297+
}
298+
got, err := r.ListAllLoadBalancerListeners(call.loadBalancerArn)
299+
assert.Equal(t, call.wantErr, err)
300+
301+
changelog, err := diff.Diff(got, call.want)
302+
assert.Nil(t, err)
303+
if len(changelog) > 0 {
304+
for _, change := range changelog {
305+
t.Errorf("%s: %v -> %v", strings.Join(change.Path, "."), change.From, change.To)
306+
}
307+
t.Fail()
238308
}
239-
t.Fail()
240309
}
241310
})
242311
}

0 commit comments

Comments
 (0)