Merge pull request aws#888 from jeskew/fix/correctly-requeue-write-re…

…quest-batch-items Correctly requeue throttled WriteRequestBatch items
rohjay · Jan 29, 2016 · 553e7ef · 553e7ef
2 parents 752b58b + 6644791
commit 553e7ef
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/src/DynamoDb/WriteRequestBatch.php b/src/DynamoDb/WriteRequestBatch.php
@@ -227,9 +227,7 @@ private function retryUnprocessed(array $unprocessed)
             foreach ($requests as $request) {
                 $this->queue[] = [
                     'table' => $table,
-                    'data'  => isset($request['PutRequest'])
-                        ? $request['PutRequest']
-                        : $request['DeleteRequest']
+                    'data'  => $request,
                 ];
             }
         }

diff --git a/tests/DynamoDb/WriteRequestBatchTest.php b/tests/DynamoDb/WriteRequestBatchTest.php
@@ -114,6 +114,7 @@ public function testUnprocessedItemsAreRequeued()
         $this->addMockResults($client, [
             new Result(['UnprocessedItems' => ['foo' => [
                 ['PutRequest' => ['Item' => ['id' => ['S' => 'b']]]],
+                ['DeleteRequest' => ['Key' => ['id' => ['S' => 'c']]]],
             ]]]),
             new Result([])
         ]);
@@ -122,8 +123,24 @@ public function testUnprocessedItemsAreRequeued()
 
         $batch->put(['id' => ['S' => 'a']]);
         $batch->put(['id' => ['S' => 'b']]);
-        $batch->flush();
+        $batch->delete(['id' => ['S' => 'c']]);
 
+        $batch->flush(false);
+        $this->assertEquals(
+            [
+                [
+                    'table' => 'foo',
+                    'data'  => ['PutRequest' => ['Item' => ['id' => ['S' => 'b']]]],
+                ],
+                [
+                    'table' => 'foo',
+                    'data'  => ['DeleteRequest' => ['Key' => ['id' => ['S' => 'c']]]],
+                ]
+            ],
+            $this->readAttribute($batch, 'queue')
+        );
+
+        $batch->flush();
         $this->assertCount(0, $this->readAttribute($batch, 'queue'));
     }