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 }