diff --git a/pkg/tsdb/cloudwatch/mocks/accounts_service.go b/pkg/tsdb/cloudwatch/mocks/accounts_service.go index 52f3aff63b5..d7b742d6848 100644 --- a/pkg/tsdb/cloudwatch/mocks/accounts_service.go +++ b/pkg/tsdb/cloudwatch/mocks/accounts_service.go @@ -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) diff --git a/pkg/tsdb/cloudwatch/mocks/oam_client.go b/pkg/tsdb/cloudwatch/mocks/oam_client.go index 2d617cc143a..932ca1579b4 100644 --- a/pkg/tsdb/cloudwatch/mocks/oam_client.go +++ b/pkg/tsdb/cloudwatch/mocks/oam_client.go @@ -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) } diff --git a/pkg/tsdb/cloudwatch/models/api.go b/pkg/tsdb/cloudwatch/models/api.go index 42d54c84e81..dc540f54d6f 100644 --- a/pkg/tsdb/cloudwatch/models/api.go +++ b/pkg/tsdb/cloudwatch/models/api.go @@ -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 { diff --git a/pkg/tsdb/cloudwatch/routes/accounts.go b/pkg/tsdb/cloudwatch/routes/accounts.go index 3c18591b55f..565a8c67ac4 100644 --- a/pkg/tsdb/cloudwatch/routes/accounts.go +++ b/pkg/tsdb/cloudwatch/routes/accounts.go @@ -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 { diff --git a/pkg/tsdb/cloudwatch/services/accounts.go b/pkg/tsdb/cloudwatch/services/accounts.go index 7f30f038050..f05b18335a3 100644 --- a/pkg/tsdb/cloudwatch/services/accounts.go +++ b/pkg/tsdb/cloudwatch/services/accounts.go @@ -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, }) diff --git a/pkg/tsdb/cloudwatch/services/accounts_test.go b/pkg/tsdb/cloudwatch/services/accounts_test.go index 7e74e6fbec9..dbbadaf660c 100644 --- a/pkg/tsdb/cloudwatch/services/accounts_test.go +++ b/pkg/tsdb/cloudwatch/services/accounts_test.go @@ -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)