From c2de3014441312ad14d35a8d034953857bc1d88b Mon Sep 17 00:00:00 2001
From: Cory Bradshaw <cobrads@amazon.com>
Date: Fri, 29 Nov 2019 14:21:51 -0800
Subject: [PATCH] CR Feedback

---
 .../lib/worker/KinesisClientLibConfiguration.java        | 9 ++++++++-
 .../services/kinesis/leases/impl/LeaseManager.java       | 2 ++
 .../kinesis/leases/impl/LeaseManagerIntegrationTest.java | 3 ++-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfiguration.java b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfiguration.java
index 9a150015..7ee2d449 100644
--- a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfiguration.java
+++ b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/KinesisClientLibConfiguration.java
@@ -361,6 +361,7 @@ public class KinesisClientLibConfiguration {
      *        {@link RecordProcessorCheckpointer#checkpoint(String)}
      * @param regionName The region name for the service
      * @param shutdownGraceMillis The number of milliseconds before graceful shutdown terminates forcefully
+     * @param billingMode The DDB Billing mode to set for lease table creation.
      */
     // CHECKSTYLE:IGNORE HiddenFieldCheck FOR NEXT 26 LINES
     // CHECKSTYLE:IGNORE ParameterNumber FOR NEXT 26 LINES
@@ -431,6 +432,7 @@ public class KinesisClientLibConfiguration {
      *        with a call to Amazon Kinesis before checkpointing for calls to
      *        {@link RecordProcessorCheckpointer#checkpoint(String)}
      * @param regionName The region name for the service
+     * @param billingMode The DDB Billing mode to set for lease table creation.
      */
     // CHECKSTYLE:IGNORE HiddenFieldCheck FOR NEXT 26 LINES
     // CHECKSTYLE:IGNORE ParameterNumber FOR NEXT 26 LINES
@@ -507,7 +509,7 @@ public class KinesisClientLibConfiguration {
         this.shardSyncStrategyType = DEFAULT_SHARD_SYNC_STRATEGY_TYPE;
         this.shardPrioritization = DEFAULT_SHARD_PRIORITIZATION;
         this.recordsFetcherFactory = new SimpleRecordsFetcherFactory();
-        this.billingMode=billingMode;
+        this.billingMode = billingMode;
     }
 
     /**
@@ -1162,6 +1164,11 @@ public class KinesisClientLibConfiguration {
         return this;
     }
 
+    /**
+     * The DDB Billing mode to set for lease table creation.
+     * @param billingMode - Either PAY_PER_REQUEST, or PROVISIONED; Defaults to PROVISIONED
+     * @return
+     */
     public KinesisClientLibConfiguration withBillingMode(BillingMode billingMode){
         this.billingMode = billingMode == null ? DEFAULT_DDB_BILLING_MODE : billingMode;
         return this;
diff --git a/src/main/java/com/amazonaws/services/kinesis/leases/impl/LeaseManager.java b/src/main/java/com/amazonaws/services/kinesis/leases/impl/LeaseManager.java
index 80a118ff..1c70e307 100644
--- a/src/main/java/com/amazonaws/services/kinesis/leases/impl/LeaseManager.java
+++ b/src/main/java/com/amazonaws/services/kinesis/leases/impl/LeaseManager.java
@@ -70,6 +70,7 @@ public class LeaseManager<T extends Lease> implements ILeaseManager<T> {
      * @param table leases table
      * @param dynamoDBClient DynamoDB client to use
      * @param serializer LeaseSerializer to use to convert to/from DynamoDB objects.
+     * @param billingMode The DDB Billing mode to set for lease table creation.
      */
     public LeaseManager(String table, AmazonDynamoDB dynamoDBClient, ILeaseSerializer<T> serializer, BillingMode billingMode) {
         this(table, dynamoDBClient, serializer, false, billingMode);
@@ -85,6 +86,7 @@ public class LeaseManager<T extends Lease> implements ILeaseManager<T> {
      * @param dynamoDBClient DynamoDB client to use
      * @param serializer lease serializer to use
      * @param consistentReads true if we want consistent reads for testing purposes.
+     * @param billingMode The DDB Billing mode to set for lease table creation.
      */
     public LeaseManager(String table, AmazonDynamoDB dynamoDBClient, ILeaseSerializer<T> serializer, boolean consistentReads, BillingMode billingMode) {
         verifyNotNull(table, "Table name cannot be null");
diff --git a/src/test/java/com/amazonaws/services/kinesis/leases/impl/LeaseManagerIntegrationTest.java b/src/test/java/com/amazonaws/services/kinesis/leases/impl/LeaseManagerIntegrationTest.java
index dd466eba..cede5be1 100644
--- a/src/test/java/com/amazonaws/services/kinesis/leases/impl/LeaseManagerIntegrationTest.java
+++ b/src/test/java/com/amazonaws/services/kinesis/leases/impl/LeaseManagerIntegrationTest.java
@@ -18,6 +18,7 @@ import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import com.amazonaws.services.dynamodbv2.model.BillingMode;
 import com.amazonaws.services.kinesis.clientlibrary.lib.worker.KinesisClientLibConfiguration;
 import junit.framework.Assert;
 
@@ -251,7 +252,7 @@ public class LeaseManagerIntegrationTest extends LeaseIntegrationTest {
     @Test
     public void testWaitUntilLeaseTableExistsPayPerRequest() throws LeasingException {
         KinesisClientLeaseManager manager = new KinesisClientLeaseManager("nagl_ShardProgress_PayPerRequest", ddbClient, true,
-                KinesisClientLibConfiguration.DEFAULT_DDB_BILLING_MODE) {
+                BillingMode.PAY_PER_REQUEST) {
 
             @Override
             long sleep(long timeToSleepMillis) {