fleet/server/logging/lambda_test.go
Lucas Manuel Rodriguez 404f0d3ac0
Migrate from aws-sdk-go v1 to v2 (#30308)
#29482

[Migrate to the AWS SDK for Go
v2](https://docs.aws.amazon.com/sdk-for-go/v2/developer-guide/migrate-gosdk.html)
documents how to migrate codebases.

QA on features that use AWS SDK Go:
- Bootstrap package:
  - upload:  
  - download: 
  - cleanup: 
- Software (upload, download, installation, etc.) 
  - Cloudfront: Luckly, this feature was already using aws-sdk-go-v2.
- Carves 
- Logging:
	- Firehose 
	- Kinesis 
- Lambda  (tested result logs to a lambda function on our AWS Dogfood
account)
- Email:
	- Amazon SES TODO ⚠️ (this is what Dogfood uses and a few customers)
- We cannot easily test locally, we can use dogfood or load testing
(AWS) environments.

---

- [X] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/guides/committing-changes.md#changes-files)
for more information.
- [ ] Manual QA for all new/changed functionality
2025-06-30 17:45:39 -03:00

84 lines
2.7 KiB
Go

package logging
import (
"context"
"errors"
"testing"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/service/lambda"
"github.com/aws/aws-sdk-go-v2/service/lambda/types"
"github.com/fleetdm/fleet/v4/server/logging/mock"
"github.com/fleetdm/fleet/v4/server/test"
"github.com/go-kit/log"
"github.com/stretchr/testify/assert"
tmock "github.com/stretchr/testify/mock"
)
func makeLambdaWriterWithMock(client LambdaAPI, functionName string) *lambdaLogWriter {
return &lambdaLogWriter{
client: client,
functionName: functionName,
logger: log.NewNopLogger(),
}
}
func TestLambdaValidateFunctionError(t *testing.T) {
m := &mock.LambdaMock{}
ctx := context.Background()
m.On("Invoke", &lambda.InvokeInput{FunctionName: aws.String("foobar"), InvocationType: types.InvocationTypeDryRun}).
Return(nil, errors.New("failed"))
writer := makeLambdaWriterWithMock(m, "foobar")
err := writer.validateFunction(ctx)
assert.Error(t, err)
m.AssertExpectations(test.Quiet(t))
}
func TestLambdaValidateFunctionErrorFunction(t *testing.T) {
m := &mock.LambdaMock{}
ctx := context.Background()
m.On("Invoke", &lambda.InvokeInput{FunctionName: aws.String("foobar"), InvocationType: types.InvocationTypeDryRun}).
Return(&lambda.InvokeOutput{FunctionError: aws.String("failed")}, nil)
writer := makeLambdaWriterWithMock(m, "foobar")
err := writer.validateFunction(ctx)
assert.Error(t, err)
m.AssertExpectations(test.Quiet(t))
}
func TestLambdaValidateFunctionSuccess(t *testing.T) {
m := &mock.LambdaMock{}
ctx := context.Background()
m.On("Invoke", &lambda.InvokeInput{FunctionName: aws.String("foobar"), InvocationType: types.InvocationTypeDryRun}).
Return(&lambda.InvokeOutput{}, nil)
writer := makeLambdaWriterWithMock(m, "foobar")
err := writer.validateFunction(ctx)
assert.NoError(t, err)
m.AssertExpectations(test.Quiet(t))
}
func TestLambdaError(t *testing.T) {
m := &mock.LambdaMock{}
m.On("Invoke", tmock.MatchedBy(
func(in *lambda.InvokeInput) bool {
return *in.FunctionName == "foobar" && in.InvocationType == ""
},
)).Return(nil, errors.New("failed"))
writer := makeLambdaWriterWithMock(m, "foobar")
err := writer.Write(context.Background(), logs)
assert.Error(t, err)
m.AssertExpectations(test.Quiet(t))
}
func TestLambdaSuccess(t *testing.T) {
m := &mock.LambdaMock{}
m.On("Invoke", tmock.MatchedBy(
func(in *lambda.InvokeInput) bool {
return len(in.Payload) > 0 && *in.FunctionName == "foobar" && in.InvocationType == ""
},
)).Return(&lambda.InvokeOutput{}, nil).
Times(len(logs))
writer := makeLambdaWriterWithMock(m, "foobar")
err := writer.Write(context.Background(), logs)
assert.NoError(t, err)
m.AssertExpectations(test.Quiet(t))
}