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 }