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

     1  package writer
     2  
     3  import (
     4  	"errors"
     5  	"testing"
     6  
     7  	"github.com/Jeffail/benthos/v3/internal/batch"
     8  	"github.com/Jeffail/benthos/v3/lib/log"
     9  	"github.com/Jeffail/benthos/v3/lib/message"
    10  	"github.com/Jeffail/benthos/v3/lib/metrics"
    11  	"github.com/aws/aws-sdk-go/aws"
    12  	"github.com/aws/aws-sdk-go/service/dynamodb"
    13  	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  )
    17  
    18  type mockDynamoDB struct {
    19  	dynamodbiface.DynamoDBAPI
    20  	fn      func(*dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error)
    21  	batchFn func(*dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error)
    22  }
    23  
    24  func (m *mockDynamoDB) PutItem(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
    25  	return m.fn(input)
    26  }
    27  
    28  func (m *mockDynamoDB) BatchWriteItem(input *dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error) {
    29  	return m.batchFn(input)
    30  }
    31  
    32  func TestDynamoDBHappy(t *testing.T) {
    33  	conf := NewDynamoDBConfig()
    34  	conf.StringColumns = map[string]string{
    35  		"id":      `${!json("id")}`,
    36  		"content": `${!json("content")}`,
    37  	}
    38  	conf.Table = "FooTable"
    39  
    40  	db, err := NewDynamoDB(conf, log.Noop(), metrics.Noop())
    41  	require.NoError(t, err)
    42  
    43  	var request map[string][]*dynamodb.WriteRequest
    44  
    45  	db.client = &mockDynamoDB{
    46  		fn: func(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
    47  			t.Error("not expected")
    48  			return nil, errors.New("not implemented")
    49  		},
    50  		batchFn: func(input *dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error) {
    51  			request = input.RequestItems
    52  			return &dynamodb.BatchWriteItemOutput{}, nil
    53  		},
    54  	}
    55  
    56  	require.NoError(t, db.Write(message.New([][]byte{
    57  		[]byte(`{"id":"foo","content":"foo stuff"}`),
    58  		[]byte(`{"id":"bar","content":"bar stuff"}`),
    59  	})))
    60  
    61  	expected := map[string][]*dynamodb.WriteRequest{
    62  		"FooTable": {
    63  			&dynamodb.WriteRequest{
    64  				PutRequest: &dynamodb.PutRequest{
    65  					Item: map[string]*dynamodb.AttributeValue{
    66  						"id": {
    67  							S: aws.String("foo"),
    68  						},
    69  						"content": {
    70  							S: aws.String("foo stuff"),
    71  						},
    72  					},
    73  				},
    74  			},
    75  			&dynamodb.WriteRequest{
    76  				PutRequest: &dynamodb.PutRequest{
    77  					Item: map[string]*dynamodb.AttributeValue{
    78  						"id": {
    79  							S: aws.String("bar"),
    80  						},
    81  						"content": {
    82  							S: aws.String("bar stuff"),
    83  						},
    84  					},
    85  				},
    86  			},
    87  		},
    88  	}
    89  
    90  	assert.Equal(t, expected, request)
    91  }
    92  
    93  func TestDynamoDBSadToGood(t *testing.T) {
    94  	t.Parallel()
    95  
    96  	conf := NewDynamoDBConfig()
    97  	conf.StringColumns = map[string]string{
    98  		"id":      `${!json("id")}`,
    99  		"content": `${!json("content")}`,
   100  	}
   101  	conf.Backoff.MaxElapsedTime = "100ms"
   102  	conf.Table = "FooTable"
   103  
   104  	db, err := NewDynamoDB(conf, log.Noop(), metrics.Noop())
   105  	require.NoError(t, err)
   106  
   107  	var batchRequest []*dynamodb.WriteRequest
   108  	var requests []*dynamodb.PutItemInput
   109  
   110  	db.client = &mockDynamoDB{
   111  		fn: func(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
   112  			requests = append(requests, input)
   113  			return nil, nil
   114  		},
   115  		batchFn: func(input *dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error) {
   116  			if len(batchRequest) > 0 {
   117  				t.Error("not expected")
   118  				return nil, errors.New("not implemented")
   119  			}
   120  			if request, ok := input.RequestItems["FooTable"]; ok {
   121  				items := make([]*dynamodb.WriteRequest, len(request))
   122  				copy(items, request)
   123  				batchRequest = items
   124  			} else {
   125  				t.Error("missing FooTable")
   126  			}
   127  			return &dynamodb.BatchWriteItemOutput{}, errors.New("woop")
   128  		},
   129  	}
   130  
   131  	require.NoError(t, db.Write(message.New([][]byte{
   132  		[]byte(`{"id":"foo","content":"foo stuff"}`),
   133  		[]byte(`{"id":"bar","content":"bar stuff"}`),
   134  		[]byte(`{"id":"baz","content":"baz stuff"}`),
   135  	})))
   136  
   137  	batchExpected := []*dynamodb.WriteRequest{
   138  		{
   139  			PutRequest: &dynamodb.PutRequest{
   140  				Item: map[string]*dynamodb.AttributeValue{
   141  					"id":      {S: aws.String("foo")},
   142  					"content": {S: aws.String("foo stuff")},
   143  				},
   144  			},
   145  		},
   146  		{
   147  			PutRequest: &dynamodb.PutRequest{
   148  				Item: map[string]*dynamodb.AttributeValue{
   149  					"id":      {S: aws.String("bar")},
   150  					"content": {S: aws.String("bar stuff")},
   151  				},
   152  			},
   153  		},
   154  		{
   155  			PutRequest: &dynamodb.PutRequest{
   156  				Item: map[string]*dynamodb.AttributeValue{
   157  					"id":      {S: aws.String("baz")},
   158  					"content": {S: aws.String("baz stuff")},
   159  				},
   160  			},
   161  		},
   162  	}
   163  
   164  	assert.Equal(t, batchExpected, batchRequest)
   165  
   166  	expected := []*dynamodb.PutItemInput{
   167  		{
   168  			TableName: aws.String("FooTable"),
   169  			Item: map[string]*dynamodb.AttributeValue{
   170  				"id":      {S: aws.String("foo")},
   171  				"content": {S: aws.String("foo stuff")},
   172  			},
   173  		},
   174  		{
   175  			TableName: aws.String("FooTable"),
   176  			Item: map[string]*dynamodb.AttributeValue{
   177  				"id":      {S: aws.String("bar")},
   178  				"content": {S: aws.String("bar stuff")},
   179  			},
   180  		},
   181  		{
   182  			TableName: aws.String("FooTable"),
   183  			Item: map[string]*dynamodb.AttributeValue{
   184  				"id":      {S: aws.String("baz")},
   185  				"content": {S: aws.String("baz stuff")},
   186  			},
   187  		},
   188  	}
   189  
   190  	assert.Equal(t, expected, requests)
   191  }
   192  
   193  func TestDynamoDBSadToGoodBatch(t *testing.T) {
   194  	t.Parallel()
   195  
   196  	conf := NewDynamoDBConfig()
   197  	conf.StringColumns = map[string]string{
   198  		"id":      `${!json("id")}`,
   199  		"content": `${!json("content")}`,
   200  	}
   201  	conf.Table = "FooTable"
   202  
   203  	db, err := NewDynamoDB(conf, log.Noop(), metrics.Noop())
   204  	require.NoError(t, err)
   205  
   206  	var requests [][]*dynamodb.WriteRequest
   207  
   208  	db.client = &mockDynamoDB{
   209  		fn: func(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
   210  			t.Error("not expected")
   211  			return nil, errors.New("not implemented")
   212  		},
   213  		batchFn: func(input *dynamodb.BatchWriteItemInput) (output *dynamodb.BatchWriteItemOutput, err error) {
   214  			if len(requests) == 0 {
   215  				output = &dynamodb.BatchWriteItemOutput{
   216  					UnprocessedItems: map[string][]*dynamodb.WriteRequest{
   217  						"FooTable": {
   218  							{
   219  								PutRequest: &dynamodb.PutRequest{
   220  									Item: map[string]*dynamodb.AttributeValue{
   221  										"id":      {S: aws.String("bar")},
   222  										"content": {S: aws.String("bar stuff")},
   223  									},
   224  								},
   225  							},
   226  						},
   227  					},
   228  				}
   229  			} else {
   230  				output = &dynamodb.BatchWriteItemOutput{}
   231  			}
   232  			if request, ok := input.RequestItems["FooTable"]; ok {
   233  				items := make([]*dynamodb.WriteRequest, len(request))
   234  				copy(items, request)
   235  				requests = append(requests, items)
   236  			} else {
   237  				t.Error("missing FooTable")
   238  			}
   239  			return
   240  		},
   241  	}
   242  
   243  	require.NoError(t, db.Write(message.New([][]byte{
   244  		[]byte(`{"id":"foo","content":"foo stuff"}`),
   245  		[]byte(`{"id":"bar","content":"bar stuff"}`),
   246  		[]byte(`{"id":"baz","content":"baz stuff"}`),
   247  	})))
   248  
   249  	expected := [][]*dynamodb.WriteRequest{
   250  		{
   251  			{
   252  				PutRequest: &dynamodb.PutRequest{
   253  					Item: map[string]*dynamodb.AttributeValue{
   254  						"id":      {S: aws.String("foo")},
   255  						"content": {S: aws.String("foo stuff")},
   256  					},
   257  				},
   258  			},
   259  			{
   260  				PutRequest: &dynamodb.PutRequest{
   261  					Item: map[string]*dynamodb.AttributeValue{
   262  						"id":      {S: aws.String("bar")},
   263  						"content": {S: aws.String("bar stuff")},
   264  					},
   265  				},
   266  			},
   267  			{
   268  				PutRequest: &dynamodb.PutRequest{
   269  					Item: map[string]*dynamodb.AttributeValue{
   270  						"id":      {S: aws.String("baz")},
   271  						"content": {S: aws.String("baz stuff")},
   272  					},
   273  				},
   274  			},
   275  		},
   276  		{
   277  			{
   278  				PutRequest: &dynamodb.PutRequest{
   279  					Item: map[string]*dynamodb.AttributeValue{
   280  						"id":      {S: aws.String("bar")},
   281  						"content": {S: aws.String("bar stuff")},
   282  					},
   283  				},
   284  			},
   285  		},
   286  	}
   287  
   288  	assert.Equal(t, expected, requests)
   289  }
   290  
   291  func TestDynamoDBSad(t *testing.T) {
   292  	t.Parallel()
   293  
   294  	conf := NewDynamoDBConfig()
   295  	conf.StringColumns = map[string]string{
   296  		"id":      `${!json("id")}`,
   297  		"content": `${!json("content")}`,
   298  	}
   299  	conf.Table = "FooTable"
   300  
   301  	db, err := NewDynamoDB(conf, log.Noop(), metrics.Noop())
   302  	require.NoError(t, err)
   303  
   304  	var batchRequest []*dynamodb.WriteRequest
   305  	var requests []*dynamodb.PutItemInput
   306  
   307  	barErr := errors.New("dont like bar")
   308  
   309  	db.client = &mockDynamoDB{
   310  		fn: func(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
   311  			if len(requests) < 3 {
   312  				requests = append(requests, input)
   313  			}
   314  			if *input.Item["id"].S == "bar" {
   315  				return nil, barErr
   316  			}
   317  			return nil, nil
   318  		},
   319  		batchFn: func(input *dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error) {
   320  			if len(batchRequest) > 0 {
   321  				t.Error("not expected")
   322  				return nil, errors.New("not implemented")
   323  			}
   324  			if request, ok := input.RequestItems["FooTable"]; ok {
   325  				items := make([]*dynamodb.WriteRequest, len(request))
   326  				copy(items, request)
   327  				batchRequest = items
   328  			} else {
   329  				t.Error("missing FooTable")
   330  			}
   331  			return &dynamodb.BatchWriteItemOutput{}, errors.New("woop")
   332  		},
   333  	}
   334  
   335  	msg := message.New([][]byte{
   336  		[]byte(`{"id":"foo","content":"foo stuff"}`),
   337  		[]byte(`{"id":"bar","content":"bar stuff"}`),
   338  		[]byte(`{"id":"baz","content":"baz stuff"}`),
   339  	})
   340  
   341  	expErr := batch.NewError(msg, errors.New("woop"))
   342  	expErr.Failed(1, barErr)
   343  	require.Equal(t, expErr, db.Write(msg))
   344  
   345  	batchExpected := []*dynamodb.WriteRequest{
   346  		{
   347  			PutRequest: &dynamodb.PutRequest{
   348  				Item: map[string]*dynamodb.AttributeValue{
   349  					"id":      {S: aws.String("foo")},
   350  					"content": {S: aws.String("foo stuff")},
   351  				},
   352  			},
   353  		},
   354  		{
   355  			PutRequest: &dynamodb.PutRequest{
   356  				Item: map[string]*dynamodb.AttributeValue{
   357  					"id":      {S: aws.String("bar")},
   358  					"content": {S: aws.String("bar stuff")},
   359  				},
   360  			},
   361  		},
   362  		{
   363  			PutRequest: &dynamodb.PutRequest{
   364  				Item: map[string]*dynamodb.AttributeValue{
   365  					"id":      {S: aws.String("baz")},
   366  					"content": {S: aws.String("baz stuff")},
   367  				},
   368  			},
   369  		},
   370  	}
   371  
   372  	assert.Equal(t, batchExpected, batchRequest)
   373  
   374  	expected := []*dynamodb.PutItemInput{
   375  		{
   376  			TableName: aws.String("FooTable"),
   377  			Item: map[string]*dynamodb.AttributeValue{
   378  				"id":      {S: aws.String("foo")},
   379  				"content": {S: aws.String("foo stuff")},
   380  			},
   381  		},
   382  		{
   383  			TableName: aws.String("FooTable"),
   384  			Item: map[string]*dynamodb.AttributeValue{
   385  				"id":      {S: aws.String("bar")},
   386  				"content": {S: aws.String("bar stuff")},
   387  			},
   388  		},
   389  		{
   390  			TableName: aws.String("FooTable"),
   391  			Item: map[string]*dynamodb.AttributeValue{
   392  				"id":      {S: aws.String("baz")},
   393  				"content": {S: aws.String("baz stuff")},
   394  			},
   395  		},
   396  	}
   397  
   398  	assert.Equal(t, expected, requests)
   399  }
   400  
   401  func TestDynamoDBSadBatch(t *testing.T) {
   402  	t.Parallel()
   403  
   404  	conf := NewDynamoDBConfig()
   405  	conf.StringColumns = map[string]string{
   406  		"id":      `${!json("id")}`,
   407  		"content": `${!json("content")}`,
   408  	}
   409  	conf.Table = "FooTable"
   410  
   411  	db, err := NewDynamoDB(conf, log.Noop(), metrics.Noop())
   412  	require.NoError(t, err)
   413  
   414  	var requests [][]*dynamodb.WriteRequest
   415  
   416  	db.client = &mockDynamoDB{
   417  		fn: func(input *dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error) {
   418  			t.Error("not expected")
   419  			return nil, errors.New("not implemented")
   420  		},
   421  		batchFn: func(input *dynamodb.BatchWriteItemInput) (output *dynamodb.BatchWriteItemOutput, err error) {
   422  			output = &dynamodb.BatchWriteItemOutput{
   423  				UnprocessedItems: map[string][]*dynamodb.WriteRequest{
   424  					"FooTable": {
   425  						{
   426  							PutRequest: &dynamodb.PutRequest{
   427  								Item: map[string]*dynamodb.AttributeValue{
   428  									"id":      {S: aws.String("bar")},
   429  									"content": {S: aws.String("bar stuff")},
   430  								},
   431  							},
   432  						},
   433  					},
   434  				},
   435  			}
   436  			if len(requests) < 2 {
   437  				if request, ok := input.RequestItems["FooTable"]; ok {
   438  					items := make([]*dynamodb.WriteRequest, len(request))
   439  					copy(items, request)
   440  					requests = append(requests, items)
   441  				} else {
   442  					t.Error("missing FooTable")
   443  				}
   444  			}
   445  			return
   446  		},
   447  	}
   448  
   449  	msg := message.New([][]byte{
   450  		[]byte(`{"id":"foo","content":"foo stuff"}`),
   451  		[]byte(`{"id":"bar","content":"bar stuff"}`),
   452  		[]byte(`{"id":"baz","content":"baz stuff"}`),
   453  	})
   454  
   455  	expErr := batch.NewError(msg, errors.New("failed to set 1 items"))
   456  	expErr.Failed(1, errors.New("failed to set item"))
   457  	require.Equal(t, expErr, db.Write(msg))
   458  
   459  	expected := [][]*dynamodb.WriteRequest{
   460  		{
   461  			{
   462  				PutRequest: &dynamodb.PutRequest{
   463  					Item: map[string]*dynamodb.AttributeValue{
   464  						"id":      {S: aws.String("foo")},
   465  						"content": {S: aws.String("foo stuff")},
   466  					},
   467  				},
   468  			},
   469  			{
   470  				PutRequest: &dynamodb.PutRequest{
   471  					Item: map[string]*dynamodb.AttributeValue{
   472  						"id":      {S: aws.String("bar")},
   473  						"content": {S: aws.String("bar stuff")},
   474  					},
   475  				},
   476  			},
   477  			{
   478  				PutRequest: &dynamodb.PutRequest{
   479  					Item: map[string]*dynamodb.AttributeValue{
   480  						"id":      {S: aws.String("baz")},
   481  						"content": {S: aws.String("baz stuff")},
   482  					},
   483  				},
   484  			},
   485  		},
   486  		{
   487  			{
   488  				PutRequest: &dynamodb.PutRequest{
   489  					Item: map[string]*dynamodb.AttributeValue{
   490  						"id":      {S: aws.String("bar")},
   491  						"content": {S: aws.String("bar stuff")},
   492  					},
   493  				},
   494  			},
   495  		},
   496  	}
   497  
   498  	assert.Equal(t, expected, requests)
   499  }