CloudWatch: Use context in aws ListSinks and ListAttachedLinks (#77083)

* Use context in aws ListSinks and ListAttachedLinks

In the current way, ListSinks and ListAttachedLinks is used which doesn't
allow cancelling the request if the context changes.
Using ListSinksWithContext and ListAttachedLinksWithContext is the
preferred way. Adding context for GetAccountsForCurrentUserOrRole
is required to pass it to ListSinks method.
This commit is contained in:
Shabeeb Khalid 2023-10-27 11:49:37 +03:00 committed by GitHub
parent edd0e80ba0
commit df41378472
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 43 additions and 36 deletions

View File

@ -1,6 +1,8 @@
package mocks
import (
"context"
"github.com/grafana/grafana/pkg/tsdb/cloudwatch/models/resources"
"github.com/stretchr/testify/mock"
)
@ -9,7 +11,7 @@ type AccountsServiceMock struct {
mock.Mock
}
func (a *AccountsServiceMock) GetAccountsForCurrentUserOrRole() ([]resources.ResourceResponse[resources.Account], error) {
func (a *AccountsServiceMock) GetAccountsForCurrentUserOrRole(ctx context.Context) ([]resources.ResourceResponse[resources.Account], error) {
args := a.Called()
return args.Get(0).([]resources.ResourceResponse[resources.Account]), args.Error(1)

View File

@ -1,6 +1,9 @@
package mocks
import (
"context"
"github.com/aws/aws-sdk-go/aws/request"
"github.com/aws/aws-sdk-go/service/oam"
"github.com/stretchr/testify/mock"
)
@ -9,12 +12,12 @@ type FakeOAMClient struct {
mock.Mock
}
func (o *FakeOAMClient) ListSinks(input *oam.ListSinksInput) (*oam.ListSinksOutput, error) {
func (o *FakeOAMClient) ListSinksWithContext(ctx context.Context, input *oam.ListSinksInput, opts ...request.Option) (*oam.ListSinksOutput, error) {
args := o.Called(input)
return args.Get(0).(*oam.ListSinksOutput), args.Error(1)
}
func (o *FakeOAMClient) ListAttachedLinks(input *oam.ListAttachedLinksInput) (*oam.ListAttachedLinksOutput, error) {
func (o *FakeOAMClient) ListAttachedLinksWithContext(ctx context.Context, input *oam.ListAttachedLinksInput, opts ...request.Option) (*oam.ListAttachedLinksOutput, error) {
args := o.Called(input)
return args.Get(0).(*oam.ListAttachedLinksOutput), args.Error(1)
}

View File

@ -42,7 +42,7 @@ type LogGroupsProvider interface {
}
type AccountsProvider interface {
GetAccountsForCurrentUserOrRole() ([]resources.ResourceResponse[resources.Account], error)
GetAccountsForCurrentUserOrRole(ctx context.Context) ([]resources.ResourceResponse[resources.Account], error)
}
type RegionsAPIProvider interface {
@ -65,8 +65,8 @@ type CloudWatchLogsAPIProvider interface {
}
type OAMAPIProvider interface {
ListSinks(*oam.ListSinksInput) (*oam.ListSinksOutput, error)
ListAttachedLinks(*oam.ListAttachedLinksInput) (*oam.ListAttachedLinksOutput, error)
ListSinksWithContext(ctx context.Context, in *oam.ListSinksInput, opts ...request.Option) (*oam.ListSinksOutput, error)
ListAttachedLinksWithContext(ctx context.Context, in *oam.ListAttachedLinksInput, opts ...request.Option) (*oam.ListAttachedLinksOutput, error)
}
type EC2APIProvider interface {

View File

@ -24,7 +24,7 @@ func AccountsHandler(ctx context.Context, pluginCtx backend.PluginContext, reqCt
return nil, models.NewHttpError("error in AccountsHandler", http.StatusInternalServerError, err)
}
accounts, err := service.GetAccountsForCurrentUserOrRole()
accounts, err := service.GetAccountsForCurrentUserOrRole(ctx)
if err != nil {
msg := "error getting accounts for current user or role"
switch {

View File

@ -1,6 +1,7 @@
package services
import (
"context"
"errors"
"fmt"
@ -20,11 +21,11 @@ func NewAccountsService(oamClient models.OAMAPIProvider) models.AccountsProvider
return &AccountsService{oamClient}
}
func (a *AccountsService) GetAccountsForCurrentUserOrRole() ([]resources.ResourceResponse[resources.Account], error) {
func (a *AccountsService) GetAccountsForCurrentUserOrRole(ctx context.Context) ([]resources.ResourceResponse[resources.Account], error) {
var nextToken *string
sinks := []*oam.ListSinksItem{}
for {
response, err := a.ListSinks(&oam.ListSinksInput{NextToken: nextToken})
response, err := a.ListSinksWithContext(ctx, &oam.ListSinksInput{NextToken: nextToken})
if err != nil {
var aerr awserr.Error
if errors.As(err, &aerr) {
@ -61,7 +62,7 @@ func (a *AccountsService) GetAccountsForCurrentUserOrRole() ([]resources.Resourc
nextToken = nil
for {
links, err := a.ListAttachedLinks(&oam.ListAttachedLinksInput{
links, err := a.ListAttachedLinksWithContext(ctx, &oam.ListAttachedLinksInput{
SinkIdentifier: sinkIdentifier,
NextToken: nextToken,
})

View File

@ -1,6 +1,7 @@
package services
import (
"context"
"fmt"
"testing"
@ -17,11 +18,11 @@ import (
func TestHandleGetAccounts(t *testing.T) {
t.Run("Should return an error in case of insufficient permissions from ListSinks", func(t *testing.T) {
fakeOAMClient := &mocks.FakeOAMClient{}
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{}, awserr.New("AccessDeniedException",
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{}, awserr.New("AccessDeniedException",
"AWS message", nil))
accounts := NewAccountsService(fakeOAMClient)
resp, err := accounts.GetAccountsForCurrentUserOrRole()
resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background())
assert.Error(t, err)
assert.Nil(t, resp)
@ -31,10 +32,10 @@ func TestHandleGetAccounts(t *testing.T) {
t.Run("Should return an error in case of any error from ListSinks", func(t *testing.T) {
fakeOAMClient := &mocks.FakeOAMClient{}
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{}, fmt.Errorf("some error"))
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{}, fmt.Errorf("some error"))
accounts := NewAccountsService(fakeOAMClient)
resp, err := accounts.GetAccountsForCurrentUserOrRole()
resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background())
assert.Error(t, err)
assert.Nil(t, resp)
@ -43,10 +44,10 @@ func TestHandleGetAccounts(t *testing.T) {
t.Run("Should return empty array in case no monitoring account exists", func(t *testing.T) {
fakeOAMClient := &mocks.FakeOAMClient{}
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{}, nil)
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{}, nil)
accounts := NewAccountsService(fakeOAMClient)
resp, err := accounts.GetAccountsForCurrentUserOrRole()
resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background())
assert.NoError(t, err)
assert.Empty(t, resp)
@ -54,26 +55,26 @@ func TestHandleGetAccounts(t *testing.T) {
t.Run("Should return one monitoring account (the first) even though ListSinks returns multiple sinks", func(t *testing.T) {
fakeOAMClient := &mocks.FakeOAMClient{}
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{
Items: []*oam.ListSinksItem{
{Name: aws.String("Account 1"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1")},
{Name: aws.String("Account 2"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group2")},
},
NextToken: new(string),
}, nil).Once()
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{
Items: []*oam.ListSinksItem{
{Name: aws.String("Account 3"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group3")},
},
NextToken: nil,
}, nil)
fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, nil)
fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, nil)
accounts := NewAccountsService(fakeOAMClient)
resp, err := accounts.GetAccountsForCurrentUserOrRole()
resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background())
assert.NoError(t, err)
fakeOAMClient.AssertNumberOfCalls(t, "ListSinks", 2)
fakeOAMClient.AssertNumberOfCalls(t, "ListSinksWithContext", 2)
require.Len(t, resp, 1)
assert.True(t, resp[0].Value.IsMonitoringAccount)
assert.Equal(t, "Account 1", resp[0].Value.Label)
@ -82,27 +83,27 @@ func TestHandleGetAccounts(t *testing.T) {
t.Run("Should merge the first sink with attached links", func(t *testing.T) {
fakeOAMClient := &mocks.FakeOAMClient{}
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{
Items: []*oam.ListSinksItem{
{Name: aws.String("Account 1"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1")},
{Name: aws.String("Account 2"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group2")},
},
NextToken: new(string),
}, nil).Once()
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{
Items: []*oam.ListSinksItem{
{Name: aws.String("Account 3"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group3")},
},
NextToken: nil,
}, nil)
fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{
fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{
Items: []*oam.ListAttachedLinksItem{
{Label: aws.String("Account 10"), LinkArn: aws.String("arn:aws:logs:us-east-1:123456789013:log-group:my-log-group10")},
{Label: aws.String("Account 11"), LinkArn: aws.String("arn:aws:logs:us-east-1:123456789014:log-group:my-log-group11")},
},
NextToken: new(string),
}, nil).Once()
fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{
fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{
Items: []*oam.ListAttachedLinksItem{
{Label: aws.String("Account 12"), LinkArn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group12")},
},
@ -110,11 +111,11 @@ func TestHandleGetAccounts(t *testing.T) {
}, nil)
accounts := NewAccountsService(fakeOAMClient)
resp, err := accounts.GetAccountsForCurrentUserOrRole()
resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background())
assert.NoError(t, err)
fakeOAMClient.AssertNumberOfCalls(t, "ListSinks", 2)
fakeOAMClient.AssertNumberOfCalls(t, "ListAttachedLinks", 2)
fakeOAMClient.AssertNumberOfCalls(t, "ListSinksWithContext", 2)
fakeOAMClient.AssertNumberOfCalls(t, "ListAttachedLinksWithContext", 2)
expectedAccounts := []resources.ResourceResponse[resources.Account]{
{Value: resources.Account{Id: "123456789012", Label: "Account 1", Arn: "arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1", IsMonitoringAccount: true}},
{Value: resources.Account{Id: "123456789013", Label: "Account 10", Arn: "arn:aws:logs:us-east-1:123456789013:log-group:my-log-group10", IsMonitoringAccount: false}},
@ -126,37 +127,37 @@ func TestHandleGetAccounts(t *testing.T) {
t.Run("Should call ListAttachedLinks with arn of first sink", func(t *testing.T) {
fakeOAMClient := &mocks.FakeOAMClient{}
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{
Items: []*oam.ListSinksItem{
{Name: aws.String("Account 1"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1")},
},
NextToken: new(string),
}, nil).Once()
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{
Items: []*oam.ListSinksItem{
{Name: aws.String("Account 3"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group3")},
},
NextToken: nil,
}, nil).Once()
fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, nil)
fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, nil)
accounts := NewAccountsService(fakeOAMClient)
_, _ = accounts.GetAccountsForCurrentUserOrRole()
_, _ = accounts.GetAccountsForCurrentUserOrRole(context.Background())
fakeOAMClient.AssertCalled(t, "ListAttachedLinks", &oam.ListAttachedLinksInput{
fakeOAMClient.AssertCalled(t, "ListAttachedLinksWithContext", &oam.ListAttachedLinksInput{
SinkIdentifier: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1"),
})
})
t.Run("Should return an error in case of any error from ListAttachedLinks", func(t *testing.T) {
fakeOAMClient := &mocks.FakeOAMClient{}
fakeOAMClient.On("ListSinks", mock.Anything).Return(&oam.ListSinksOutput{
fakeOAMClient.On("ListSinksWithContext", mock.Anything).Return(&oam.ListSinksOutput{
Items: []*oam.ListSinksItem{{Name: aws.String("Account 1"), Arn: aws.String("arn:aws:logs:us-east-1:123456789012:log-group:my-log-group1")}},
}, nil)
fakeOAMClient.On("ListAttachedLinks", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, fmt.Errorf("some error")).Once()
fakeOAMClient.On("ListAttachedLinksWithContext", mock.Anything).Return(&oam.ListAttachedLinksOutput{}, fmt.Errorf("some error")).Once()
accounts := NewAccountsService(fakeOAMClient)
resp, err := accounts.GetAccountsForCurrentUserOrRole()
resp, err := accounts.GetAccountsForCurrentUserOrRole(context.Background())
assert.Error(t, err)
assert.Nil(t, resp)