github.com/Jeffail/benthos/v3@v3.65.0/lib/output/writer/kinesis_integration_test.go (about)

     1  package writer
     2  
     3  import (
     4  	"bytes"
     5  	"flag"
     6  	"fmt"
     7  	"regexp"
     8  	"strconv"
     9  	"strings"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/Jeffail/benthos/v3/lib/log"
    14  	"github.com/Jeffail/benthos/v3/lib/message"
    15  	"github.com/Jeffail/benthos/v3/lib/metrics"
    16  	sess "github.com/Jeffail/benthos/v3/lib/util/aws/session"
    17  	"github.com/aws/aws-sdk-go/aws"
    18  	"github.com/aws/aws-sdk-go/aws/credentials"
    19  	"github.com/aws/aws-sdk-go/aws/session"
    20  	"github.com/aws/aws-sdk-go/service/kinesis"
    21  	"github.com/ory/dockertest/v3"
    22  )
    23  
    24  func TestKinesisIntegration(t *testing.T) {
    25  	if m := flag.Lookup("test.run").Value.String(); m == "" || regexp.MustCompile(strings.Split(m, "/")[0]).FindString(t.Name()) == "" {
    26  		t.Skip("Skipping as execution was not requested explicitly using go test -run ^TestIntegration$")
    27  	}
    28  
    29  	if testing.Short() {
    30  		t.Skip("Skipping integration test in short mode")
    31  	}
    32  
    33  	pool, err := dockertest.NewPool("")
    34  	if err != nil {
    35  		t.Skipf("Could not connect to docker: %s", err)
    36  	}
    37  	pool.MaxWait = time.Second * 30
    38  
    39  	// start mysql container with binlog enabled
    40  	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
    41  		Repository: "vsouza/kinesis-local",
    42  		Cmd: []string{
    43  			"--createStreamMs=5",
    44  		},
    45  	})
    46  	if err != nil {
    47  		t.Fatalf("Could not start resource: %v", err)
    48  	}
    49  	defer func() {
    50  		if err := pool.Purge(resource); err != nil {
    51  			t.Logf("Failed to clean up docker resource: %v", err)
    52  		}
    53  	}()
    54  
    55  	port, err := strconv.ParseInt(resource.GetPort("4567/tcp"), 10, 64)
    56  	if err != nil {
    57  		t.Fatal(err)
    58  	}
    59  
    60  	endpoint := fmt.Sprintf("http://localhost:%d", port)
    61  	config := KinesisConfig{
    62  		Stream:       "foo",
    63  		PartitionKey: "${!json(\"id\")}",
    64  	}
    65  	config.Region = "us-east-1"
    66  	config.Endpoint = endpoint
    67  	config.Credentials = sess.CredentialsConfig{
    68  		ID:     "xxxxxx",
    69  		Secret: "xxxxxx",
    70  		Token:  "xxxxxx",
    71  	}
    72  
    73  	// bootstrap kinesis
    74  	client := kinesis.New(session.Must(session.NewSession(&aws.Config{
    75  		Credentials: credentials.NewStaticCredentials("xxxxx", "xxxxx", "xxxxx"),
    76  		Endpoint:    aws.String(endpoint),
    77  		Region:      aws.String("us-east-1"),
    78  	})))
    79  	if err := pool.Retry(func() error {
    80  		_, err := client.CreateStream(&kinesis.CreateStreamInput{
    81  			ShardCount: aws.Int64(1),
    82  			StreamName: aws.String(config.Stream),
    83  		})
    84  		return err
    85  	}); err != nil {
    86  		t.Fatalf("Could not connect to docker resource: %s", err)
    87  	}
    88  
    89  	t.Run("testKinesisConnect", func(t *testing.T) {
    90  		testKinesisConnect(t, config, client)
    91  	})
    92  }
    93  
    94  func testKinesisConnect(t *testing.T, c KinesisConfig, client *kinesis.Kinesis) {
    95  	met := metrics.Noop()
    96  	log := log.Noop()
    97  	r, err := NewKinesis(c, log, met)
    98  	if err != nil {
    99  		t.Fatal(err)
   100  	}
   101  
   102  	if err := r.Connect(); err != nil {
   103  		t.Fatal(err)
   104  	}
   105  	defer func() {
   106  		r.CloseAsync()
   107  		if err := r.WaitForClose(time.Millisecond); err != nil {
   108  			t.Error(err)
   109  		}
   110  	}()
   111  
   112  	records := [][]byte{
   113  		[]byte(`{"foo":"bar","id":123}`),
   114  		[]byte(`{"foo":"baz","id":456}`),
   115  		[]byte(`{"foo":"qux","id":789}`),
   116  	}
   117  
   118  	msg := message.New(nil)
   119  	for _, record := range records {
   120  		msg.Append(message.NewPart(record))
   121  	}
   122  
   123  	if err := r.Write(msg); err != nil {
   124  		t.Fatal(err)
   125  	}
   126  
   127  	iterator, err := client.GetShardIterator(&kinesis.GetShardIteratorInput{
   128  		ShardId:           aws.String("shardId-000000000000"),
   129  		ShardIteratorType: aws.String(kinesis.ShardIteratorTypeTrimHorizon),
   130  		StreamName:        aws.String(c.Stream),
   131  	})
   132  	if err != nil {
   133  		t.Fatal(err)
   134  	}
   135  
   136  	out, err := client.GetRecords(&kinesis.GetRecordsInput{
   137  		Limit:         aws.Int64(10),
   138  		ShardIterator: iterator.ShardIterator,
   139  	})
   140  	if err != nil {
   141  		t.Error(err)
   142  	}
   143  	if act, exp := len(out.Records), len(records); act != exp {
   144  		t.Fatalf("Expected GetRecords response to have records with length of %d, got %d", exp, act)
   145  	}
   146  	for i, record := range records {
   147  		if !bytes.Equal(out.Records[i].Data, record) {
   148  			t.Errorf("Expected record %d to equal %v, got %v", i, record, out.Records[i])
   149  		}
   150  	}
   151  }