fix BatchGet request grouping (#240)

guregu · Jul 3, 2024 · 62e3395 · 62e3395
1 parent 58b8919
commit 62e3395
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/batch_test.go b/batch_test.go
@@ -188,3 +188,14 @@ func TestBatchEmptyInput(t *testing.T) {
 		t.Error("unexpected error", err)
 	}
 }
+
+func TestBatchGetInputSize(t *testing.T) {
+	// see: https://github.com/guregu/dynamo/issues/240
+	table := testDB.Table(testTableWidgets)
+	keys := []Keyed{Keys{1, "abc"}, Keys{2, "def"}}
+	batch := table.Batch("UserID", "Time").Get(keys...)
+	input := batch.input(0)
+	if size := len(input.RequestItems[testTableWidgets].Keys); size != len(keys) {
+		t.Error("input amalgamation size mismatch. want:", len(keys), "got:", size)
+	}
+}
diff --git a/batchget.go b/batchget.go
@@ -232,6 +232,7 @@ func (bg *BatchGet) input(start int) *dynamodb.BatchGetItemInput {
 			continue
 		}
 		kas.Keys = append(kas.Keys, get.keys())
+		in.RequestItems[table] = kas
 	}
 	return in
 }