From 87677255e13166b1a4db01d74885f5d68d883dde Mon Sep 17 00:00:00 2001
From: Cai41 <caiyang0226@outlook.com>
Date: Fri, 5 Jun 2020 15:45:22 -0700
Subject: [PATCH] DDB is considered healthy in UPDATING status

---
 .../dynamodb/DynamoDBLeaseRefresher.java      |  3 +-
 .../dynamodb/DynamoDBLeaseRefresherTest.java  | 44 +++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresher.java b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresher.java
index 13a531c11..10526ea65 100644
--- a/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresher.java
+++ b/amazon-kinesis-client/src/main/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresher.java
@@ -190,7 +190,8 @@ public boolean createLeaseTableIfNotExists(@NonNull final Long readCapacity, @No
      */
     @Override
     public boolean leaseTableExists() throws DependencyException {
-        return TableStatus.ACTIVE == tableStatus();
+        TableStatus tableStatus = tableStatus();
+        return TableStatus.ACTIVE == tableStatus || TableStatus.UPDATING == tableStatus;
     }
 
     private TableStatus tableStatus() throws DependencyException {
diff --git a/amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresherTest.java b/amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresherTest.java
index 3e3143479..5188c41b3 100644
--- a/amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresherTest.java
+++ b/amazon-kinesis-client/src/test/java/software/amazon/kinesis/leases/dynamodb/DynamoDBLeaseRefresherTest.java
@@ -15,6 +15,8 @@
 package software.amazon.kinesis.leases.dynamodb;
 
 import static org.hamcrest.CoreMatchers.equalTo;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyBoolean;
 import static org.mockito.Matchers.anyLong;
@@ -54,6 +56,8 @@
 import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException;
 import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
 import software.amazon.awssdk.services.dynamodb.model.ScanResponse;
+import software.amazon.awssdk.services.dynamodb.model.TableDescription;
+import software.amazon.awssdk.services.dynamodb.model.TableStatus;
 import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.UpdateItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.BillingMode;
@@ -149,6 +153,46 @@ public void testCreateLeaseIfNotExistsTimesOut() throws Exception {
         verifyCancel(mockPutItemFuture, () -> leaseRefresher.createLeaseIfNotExists(lease));
     }
 
+    @Test
+    public void testWaitUntilLeaseTableExistsUpdatingStatus() throws Exception {
+        when(dynamoDbClient.describeTable(any(DescribeTableRequest.class))).thenReturn(mockDescribeTableFuture);
+        when(mockDescribeTableFuture.get(anyLong(), any()))
+                .thenReturn(DescribeTableResponse.builder()
+                        .table(TableDescription.builder().tableStatus(TableStatus.UPDATING).build())
+                        .build());
+        assertTrue(leaseRefresher.waitUntilLeaseTableExists(0, 0));
+    }
+
+    @Test
+    public void testWaitUntilLeaseTableExistsActiveStatus() throws Exception {
+        when(dynamoDbClient.describeTable(any(DescribeTableRequest.class))).thenReturn(mockDescribeTableFuture);
+        when(mockDescribeTableFuture.get(anyLong(), any()))
+                .thenReturn(DescribeTableResponse.builder()
+                        .table(TableDescription.builder().tableStatus(TableStatus.ACTIVE).build())
+                        .build());
+        assertTrue(leaseRefresher.waitUntilLeaseTableExists(0, 0));
+    }
+
+    @Test
+    public void testWaitUntilLeaseTableExistsCreatingStatus() throws Exception {
+        when(dynamoDbClient.describeTable(any(DescribeTableRequest.class))).thenReturn(mockDescribeTableFuture);
+        when(mockDescribeTableFuture.get(anyLong(), any()))
+                .thenReturn(DescribeTableResponse.builder()
+                        .table(TableDescription.builder().tableStatus(TableStatus.CREATING).build())
+                        .build());
+        assertFalse(leaseRefresher.waitUntilLeaseTableExists(0, 0));
+    }
+
+    @Test
+    public void testWaitUntilLeaseTableExistsDeletingStatus() throws Exception {
+        when(dynamoDbClient.describeTable(any(DescribeTableRequest.class))).thenReturn(mockDescribeTableFuture);
+        when(mockDescribeTableFuture.get(anyLong(), any()))
+                .thenReturn(DescribeTableResponse.builder()
+                        .table(TableDescription.builder().tableStatus(TableStatus.DELETING).build())
+                        .build());
+        assertFalse(leaseRefresher.waitUntilLeaseTableExists(0, 0));
+    }
+
     @Test
     public void testGetLeaseTimesOut() throws Exception {
         TimeoutException te = setRuleForDependencyTimeout();