github.com/Jeffail/benthos/v3@v3.65.0/lib/test/integration/aws_kinesis_test.go (about)

     1  package integration
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/Jeffail/benthos/v3/internal/integration"
    10  	"github.com/aws/aws-sdk-go/aws"
    11  	"github.com/aws/aws-sdk-go/aws/credentials"
    12  	"github.com/aws/aws-sdk-go/aws/session"
    13  	"github.com/aws/aws-sdk-go/service/kinesis"
    14  	"github.com/ory/dockertest/v3"
    15  	"github.com/stretchr/testify/assert"
    16  	"github.com/stretchr/testify/require"
    17  )
    18  
    19  func createKinesisShards(ctx context.Context, awsPort, id string, numShards int) error {
    20  	endpoint := fmt.Sprintf("http://localhost:%v", awsPort)
    21  
    22  	client := kinesis.New(session.Must(session.NewSession(&aws.Config{
    23  		Credentials: credentials.NewStaticCredentials("xxxxx", "xxxxx", "xxxxx"),
    24  		Endpoint:    aws.String(endpoint),
    25  		Region:      aws.String("us-east-1"),
    26  	})))
    27  
    28  	_, err := client.CreateStreamWithContext(ctx, &kinesis.CreateStreamInput{
    29  		ShardCount: aws.Int64(int64(numShards)),
    30  		StreamName: aws.String("stream-" + id),
    31  	})
    32  	if err != nil {
    33  		return err
    34  	}
    35  
    36  	// wait for stream to exist
    37  	return client.WaitUntilStreamExistsWithContext(ctx, &kinesis.DescribeStreamInput{
    38  		StreamName: aws.String("stream-" + id),
    39  	})
    40  }
    41  
    42  var _ = registerIntegrationTest("aws_kinesis", func(t *testing.T) {
    43  	t.Parallel()
    44  
    45  	pool, err := dockertest.NewPool("")
    46  	require.NoError(t, err)
    47  
    48  	pool.MaxWait = time.Second * 30
    49  
    50  	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
    51  		Repository:   "localstack/localstack",
    52  		Tag:          "0.12.12", // Looks like we're bitten by https://github.com/localstack/localstack/issues/4522
    53  		ExposedPorts: []string{"4566/tcp"},
    54  		Env:          []string{"SERVICES=dynamodb,kinesis"},
    55  	})
    56  	require.NoError(t, err)
    57  	t.Cleanup(func() {
    58  		assert.NoError(t, pool.Purge(resource))
    59  	})
    60  
    61  	resource.Expire(900)
    62  
    63  	require.NoError(t, pool.Retry(func() error {
    64  		return createKinesisShards(context.Background(), resource.GetPort("4566/tcp"), "testtable", 2)
    65  	}))
    66  
    67  	template := `
    68  output:
    69    aws_kinesis:
    70      endpoint: http://localhost:$PORT
    71      region: us-east-1
    72      stream: stream-$ID
    73      partition_key: ${! uuid_v4() }
    74      max_in_flight: $MAX_IN_FLIGHT
    75      credentials:
    76        id: xxxxx
    77        secret: xxxxx
    78        token: xxxxx
    79      batching:
    80        count: $OUTPUT_BATCH_COUNT
    81  
    82  input:
    83    aws_kinesis:
    84      endpoint: http://localhost:$PORT
    85      streams: [ stream-$ID$VAR1 ]
    86      checkpoint_limit: $VAR2
    87      dynamodb:
    88        table: stream-$ID
    89        create: true
    90      start_from_oldest: true
    91      region: us-east-1
    92      credentials:
    93        id: xxxxx
    94        secret: xxxxx
    95        token: xxxxx
    96  `
    97  
    98  	suite := integration.StreamTests(
    99  		integration.StreamTestOpenClose(),
   100  		integration.StreamTestSendBatch(10),
   101  		integration.StreamTestSendBatchCount(10),
   102  		integration.StreamTestStreamSequential(200),
   103  		integration.StreamTestStreamParallel(200),
   104  		integration.StreamTestStreamParallelLossy(200),
   105  		integration.StreamTestStreamParallelLossyThroughReconnect(200),
   106  	)
   107  
   108  	t.Run("with static shards", func(t *testing.T) {
   109  		suite.Run(
   110  			t, template,
   111  			integration.StreamTestOptPreTest(func(t testing.TB, ctx context.Context, testID string, vars *integration.StreamTestConfigVars) {
   112  				streamName := "stream-" + testID
   113  				vars.Var1 = fmt.Sprintf(":0,%v:1", streamName)
   114  				require.NoError(t, createKinesisShards(ctx, resource.GetPort("4566/tcp"), testID, 2))
   115  			}),
   116  			integration.StreamTestOptPort(resource.GetPort("4566/tcp")),
   117  			integration.StreamTestOptAllowDupes(),
   118  			integration.StreamTestOptVarTwo("10"),
   119  		)
   120  	})
   121  
   122  	t.Run("with balanced shards", func(t *testing.T) {
   123  		suite.Run(
   124  			t, template,
   125  			integration.StreamTestOptPreTest(func(t testing.TB, ctx context.Context, testID string, vars *integration.StreamTestConfigVars) {
   126  				require.NoError(t, createKinesisShards(ctx, resource.GetPort("4566/tcp"), testID, 2))
   127  			}),
   128  			integration.StreamTestOptPort(resource.GetPort("4566/tcp")),
   129  			integration.StreamTestOptAllowDupes(),
   130  			integration.StreamTestOptVarTwo("10"),
   131  		)
   132  	})
   133  
   134  	t.Run("single shard", func(t *testing.T) {
   135  		integration.StreamTests(
   136  			integration.StreamTestCheckpointCapture(),
   137  		).Run(
   138  			t, template,
   139  			integration.StreamTestOptPreTest(func(t testing.TB, ctx context.Context, testID string, vars *integration.StreamTestConfigVars) {
   140  				require.NoError(t, createKinesisShards(ctx, resource.GetPort("4566/tcp"), testID, 1))
   141  			}),
   142  			integration.StreamTestOptPort(resource.GetPort("4566/tcp")),
   143  			integration.StreamTestOptAllowDupes(),
   144  			integration.StreamTestOptVarOne(":0"),
   145  			integration.StreamTestOptVarTwo("10"),
   146  		)
   147  	})
   148  })