@Test
public void testCreate() {
    AmazonDynamoDB proxyClient = getProxyClient();
    AmazonDynamoDB awsClient = getAwsClient();

    CreateTableRequest req = new CreateTableRequest()
            .withTableName("foo")
            .withProvisionedThroughput(new ProvisionedThroughput().withReadCapacityUnits(100L).withWriteCapacityUnits(100L))
            .withKeySchema(new KeySchemaElement("Name", KeyType.HASH))
            .withAttributeDefinitions(new AttributeDefinition("Name", ScalarAttributeType.S));

    proxyClient.createTable(req);
    awsClient.createTable(req);

    DescribeTableResult r = proxyClient.describeTable("foo");
    DescribeTableResult r2 = proxyClient.describeTable("foo");

    Date now = new Date();
    r.getTable().withCreationDateTime(now);
    r2.getTable().withCreationDateTime(now);

    Assert.assertEquals(r, r2);
}
 

@Test
public void testCreateTableWithWait() throws Exception {
    // Create fake responses from AWS. First response is still creating the table, second response the table
    // has become active.
    TableDescription creatingDescription = constructTableDescription(TableStatus.CREATING);
    TableDescription createdDescription = constructTableDescription(TableStatus.ACTIVE);
    CreateTableResult mockCreateResult = new CreateTableResult().withTableDescription(creatingDescription);
    DescribeTableResult mockDescribeResultCreating = new DescribeTableResult().withTable(creatingDescription);
    DescribeTableResult mockDescribeResultCreated = new DescribeTableResult().withTable(createdDescription);

    // Create the table.
    CreateTableRequest expectedRequest = dynamoDB.constructCreateTableRequest();
    when(mockDynamoDBClient.createTable(expectedRequest)).thenReturn(mockCreateResult);
    when(mockDynamoDBClient.describeTable(tableName)).thenReturn(mockDescribeResultCreating, mockDescribeResultCreated);
    assertEquals(dynamoDB.create(), TEST_ARN);

    verify(mockDynamoDBClient, times(1)).createTable(expectedRequest);
    verify(mockDynamoDBClient, times(2)).describeTable(tableName);
}
 

public TableDescription describeTable(String tableName) {
  final DescribeTableRequest describeTablesRequest = new DescribeTableRequest()
      .withTableName(tableName);
  try {
    RetryResult<DescribeTableResult> describeResult = getRetryDriver().runWithRetry(
        new Callable<DescribeTableResult>() {
          @Override
          public DescribeTableResult call() {
            DescribeTableResult result = dynamoDB.describeTable(describeTablesRequest);
            log.info("Describe table output: " + result);
            return result;
          }
        }, null, null);
    return describeResult.result.getTable();
  } catch (Exception e) {
    throw new RuntimeException("Could not lookup table " + tableName + " in DynamoDB.", e);
  }
}
 

private static void awaitTableCreation(String tableName) {
    Integer retries = 0;
    Boolean created = false;
    while (!created && retries < 100) {
        DescribeTableResult result = StreamsAdapterDemoHelper.describeTable(dynamoDBClient, tableName);
        created = result.getTable().getTableStatus().equals("ACTIVE");
        if (created) {
            System.out.println("Table is active.");
            return;
        }
        else {
            retries++;
            try {
                Thread.sleep(1000);
            }
            catch (InterruptedException e) {
                // do nothing
            }
        }
    }
    System.out.println("Timeout after table creation. Exiting...");
    cleanupAndExit(1);
}
 

@Test
public void test_updateTable() throws Exception {
  createTable();

  DescribeTableResult describeResult = dynamoDb.describeTable(TEST_TABLE_NAME);
  Long readUnits = describeResult.getTable().getProvisionedThroughput().getReadCapacityUnits();

  assertThat(readUnits, equalTo(UNITS));

  UpdateTableResult updateResult = dynamoDb.updateTable(TEST_TABLE_NAME, new ProvisionedThroughput()
    .withReadCapacityUnits(new Long(200))
    .withWriteCapacityUnits(new Long(200)));
  String tableName = updateResult.getTableDescription().getTableName();

  assertThat(tableName, equalTo(TEST_TABLE_NAME));

  ListTablesResult listResult = listTables();

  describeResult = dynamoDb.describeTable(TEST_TABLE_NAME);
  readUnits = describeResult.getTable().getProvisionedThroughput().getReadCapacityUnits();

  assertThat(readUnits, equalTo(new Long(200)));
}
 

private static void awaitTableCreation(String tableName) {
    Integer retries = 0;
    Boolean created = false;
    while(!created && retries < 100) {
        DescribeTableResult result = StreamsAdapterDemoHelper.describeTable(dynamoDBClient, tableName);
        created = result.getTable().getTableStatus().equals("ACTIVE");
        if (created) {
            System.out.println("Table is active.");
            return;
        } else {
            retries++;
            try {
                Thread.sleep(1000);
            } catch(InterruptedException e) {
                // do nothing
            }
        }
    }
    System.out.println("Timeout after table creation. Exiting...");
    cleanupAndExit(1);
}
 

public void waitForTableActive(String tableName, long waitTimeSeconds) throws InterruptedException {
    if(waitTimeSeconds < 0) {
        throw new IllegalArgumentException("Invalid waitTimeSeconds " + waitTimeSeconds);
    }
    
    long startTimeMs = System.currentTimeMillis();
    long elapsedMs = 0;
    do {
        DescribeTableResult describe = client.describeTable(new DescribeTableRequest().withTableName(tableName));
        String status = describe.getTable().getTableStatus();
        if(TableStatus.ACTIVE.toString().equals(status)) {
            return;
        }
        if(TableStatus.DELETING.toString().equals(status)) {
            throw new ResourceInUseException("Table " + tableName + " is " + status + ", and waiting for it to become ACTIVE is not useful.");
        }
        Thread.sleep(10 * 1000);
        elapsedMs = System.currentTimeMillis() - startTimeMs; 
    } while(elapsedMs / 1000.0 < waitTimeSeconds);
    
    throw new ResourceInUseException("Table " + tableName + " did not become ACTIVE after " + waitTimeSeconds + " seconds.");
}
 

public void waitForTableDeleted(String tableName, long waitTimeSeconds) throws InterruptedException {
    
    if(waitTimeSeconds < 0) {
        throw new IllegalArgumentException("Invalid waitTimeSeconds " + waitTimeSeconds);
    }
    
    long startTimeMs = System.currentTimeMillis();
    long elapsedMs = 0;
    do {
        try {
            DescribeTableResult describe = client.describeTable(new DescribeTableRequest().withTableName(tableName));
            String status = describe.getTable().getTableStatus();
            if(! TableStatus.DELETING.toString().equals(status)) {
                throw new ResourceInUseException("Table " + tableName + " is " + status + ", and waiting for it to not exist is only useful if it is DELETING.");
            }
        } catch (ResourceNotFoundException e) {
            return;
        }
        Thread.sleep(10 * 1000);
        elapsedMs = System.currentTimeMillis() - startTimeMs; 
    } while(elapsedMs / 1000.0 < waitTimeSeconds);
    
    throw new ResourceInUseException("Table " + tableName + " was not deleted after " + waitTimeSeconds + " seconds.");
}
 

/**
 * Fetch dynamo DB tables test.
 *
 * @throws Exception the exception
 */
@SuppressWarnings("static-access")
@Test
public void fetchDynamoDBTablesTest() throws Exception {
    
    mockStatic(AmazonDynamoDBClientBuilder.class);
    AmazonDynamoDB awsClient = PowerMockito.mock(AmazonDynamoDB.class);
    AmazonDynamoDBClientBuilder amazonDynamoDBClientBuilder = PowerMockito.mock(AmazonDynamoDBClientBuilder.class);
    AWSStaticCredentialsProvider awsStaticCredentialsProvider = PowerMockito.mock(AWSStaticCredentialsProvider.class);
    PowerMockito.whenNew(AWSStaticCredentialsProvider.class).withAnyArguments().thenReturn(awsStaticCredentialsProvider);
    when(amazonDynamoDBClientBuilder.standard()).thenReturn(amazonDynamoDBClientBuilder);
    when(amazonDynamoDBClientBuilder.withCredentials(anyObject())).thenReturn(amazonDynamoDBClientBuilder);
    when(amazonDynamoDBClientBuilder.withRegion(anyString())).thenReturn(amazonDynamoDBClientBuilder);
    when(amazonDynamoDBClientBuilder.build()).thenReturn(awsClient);
    
    ListTablesResult listTableResult = new ListTablesResult();
    List<String> tables = new ArrayList<>();
    tables.add(new String());
    listTableResult.setTableNames(tables);
    when(awsClient.listTables()).thenReturn(listTableResult);
    
    DescribeTableResult describeTableResult = new DescribeTableResult();
    TableDescription table = new TableDescription();
    table.setTableArn("tableArn");
    describeTableResult.setTable(table);
    when(awsClient.describeTable(anyString())).thenReturn(describeTableResult);
    
    when(awsClient.listTagsOfResource(anyObject())).thenReturn(new ListTagsOfResourceResult());
    assertThat(inventoryUtil.fetchDynamoDBTables(new BasicSessionCredentials("awsAccessKey", "awsSecretKey", "sessionToken"), 
            "skipRegions", "account","accountName").size(), is(1));
    
}
 

DescribeTableResult fetchTableDescription(Map<String, Object> properties) {
    AWSCredentials credentials = new BasicAWSCredentials(properties.get("accessKey").toString(),
        properties.get("secretKey").toString());
    AWSCredentialsProvider credentialsProvider = new AWSStaticCredentialsProvider(credentials);
    AmazonDynamoDB client = AmazonDynamoDBClientBuilder.standard().withCredentials(credentialsProvider)
        .withRegion(Regions.valueOf(properties.get("region").toString())).build();

    return client.describeTable(properties.get("tableName").toString());
}
 

@Override
protected MessageProducer createConsumerEndpoint(ConsumerDestination destination,
		String group,
		ExtendedConsumerProperties<KinesisConsumerProperties> properties) {

	ConsumerDestination destinationToUse = destination;

	if (properties.getExtension().isDynamoDbStreams()) {
		DescribeTableResult describeTableResult = this.dynamoDBClient.describeTable(destinationToUse.getName());
		String latestStreamArn = describeTableResult.getTable().getLatestStreamArn();
		if (StringUtils.hasText(latestStreamArn)) {
			destinationToUse = new KinesisConsumerDestination(latestStreamArn, Collections.emptyList());
		}
		else {
			throw new ProvisioningException("The DynamoDB table ["
					+ destinationToUse.getName()
					+ "] doesn't have Streams enabled.");
		}
	}
	else {
		this.streamsInUse.add(destinationToUse.getName());
	}

	MessageProducer adapter;
	if (this.configurationProperties.isKplKclEnabled()) {
		adapter = createKclConsumerEndpoint(destinationToUse, group, properties);
	}
	else {
		adapter = createKinesisConsumerEndpoint(destinationToUse, group, properties);
	}

	return adapter;
}
 

private void waitForTableToBecomeActive() {
    int retries = 0;
    String tableStatus = "Unknown";
    try {
        while (retries < MAX_RETRIES) {
            log.info("Waiting for table to become active...");
            Thread.sleep(SLEEP_TIME);
            DescribeTableResult result = client.describeTable(tableName);
            tableStatus = result.getTable().getTableStatus();

            if (tableStatus.equals(TableStatus.ACTIVE.toString()) ||
                    tableStatus.equals(TableStatus.UPDATING.toString())) {
                return;
            }

            if (tableStatus.equals(TableStatus.DELETING.toString())) {
                throw new UnexpectedStateException(
                        tableName, tableStatus, TableStatus.ACTIVE.toString(),
                        "Table state changed to 'DELETING' before creation was confirmed");
            }
            retries++;
        }
    } catch (InterruptedException e) {
        throw new UnexpectedStateException(tableName, tableStatus, TableStatus.ACTIVE.toString(),
                                           "Error occurred while waiting for DynamoDB table", e);
    }
    throw new UnexpectedStateException(tableName, tableStatus, TableStatus.ACTIVE.toString(),
                                       "DynamoDB table did not become active before timeout");
}
 

private Optional<TableDescription> describeTable() {
    try {
        DescribeTableResult result = client.describeTable(tableName);
        return Optional.of(result.getTable());
    } catch (ResourceNotFoundException e) {
        return Optional.empty();
    }

}
 

@Test
public void testDeleteTableWithWait() throws Exception {
    // Create fake responses from AWS.
    TableDescription deletingDescription = constructTableDescription(TableStatus.DELETING);
    DescribeTableResult mockDescribeResult = new DescribeTableResult().withTable(deletingDescription);

    // Delete the table. First response the table is still deleting, the second response the table has deleted
    // and the ResourceNotFoundException is thrown.
    when(mockDynamoDBClient.describeTable(tableName)).thenReturn(mockDescribeResult).thenThrow(
            new ResourceNotFoundException("Table not found"));
    dynamoDB.delete();

    verify(mockDynamoDBClient, times(1)).deleteTable(tableName);
    verify(mockDynamoDBClient, times(2)).describeTable(tableName);
}
 

private DescribeTableResult describeTable(final DescribeTableRequest request) throws BackendException {
    controlPlaneRateLimiter.acquire();
    final Timer.Context apiTimerContext = getTimerContext(DESCRIBE_TABLE, request.getTableName());
    DescribeTableResult result;
    try {
        result = client.describeTable(request);
    } catch (final Exception e) {
        throw processDynamoDbApiException(e, DESCRIBE_TABLE, request.getTableName());
    } finally {
        apiTimerContext.stop();
    }
    return result;
}
 

private boolean leaseTableExists() {
  DescribeTableRequest request = new DescribeTableRequest();
  request.setTableName(conf.applicationName);
  DescribeTableResult result;
  try {
    result = dynamoDBClient.describeTable(request);
  } catch (ResourceNotFoundException e) {
    LOG.debug("Lease table '{}' does not exist", conf.applicationName);
    return false;
  }

  TableStatus tableStatus = TableStatus.fromValue(result.getTable().getTableStatus());
  LOG.debug("Lease table exists and is in '{}' state", tableStatus);
  return tableStatus == TableStatus.ACTIVE;
}
 

@Test
public void test_describeTable() throws Exception {
  createTable();

  DescribeTableResult result = dynamoDb.describeTable(TEST_TABLE_NAME);
  String tableName = result.getTable().getTableName();

  assertThat(tableName, equalTo(TEST_TABLE_NAME));
}
 

/**
 * Creates DynamoDB table. If the table already exists, it validates the key schema. If the key schemas match, a
 * warning is logged, otherwise an exception is raised.
 *
 * @param dynamoDB
 *            {@link AmazonDynamoDB} used to create the table specified in the request.
 * @param request
 *            Request for creating a table.
 * @return TableDescription of the existing table or newly created table
 */
public static TableDescription createTable(final AmazonDynamoDB dynamoDB, final CreateTableRequest request) {
    try {
        final DescribeTableResult result = dynamoDB.describeTable(request.getTableName());
        if (!request.getKeySchema().equals(result.getTable().getKeySchema())) {
            throw new IllegalStateException("Table " + request.getTableName()
                + " already exists and has an invalid schema");
        }
        LOGGER.warning("Table " + request.getTableName() + " already exists");
        return result.getTable();
    } catch (final ResourceNotFoundException e) {
        return dynamoDB.createTable(request).getTableDescription();
    }
}
 

@Test
public void testCreateTableTableAlreadyExistsCorrectKeySchema() {
    final Collection<AttributeDefinition> ads = Arrays.asList(new AttributeDefinition("Hash", ScalarAttributeType.S));
    final Collection<KeySchemaElement> kses = Arrays.asList(new KeySchemaElement("Hash", KeyType.HASH));
    final TableDescription description = new TableDescription().withAttributeDefinitions(ads).withKeySchema(kses)
        .withTableName(tableName);
    final DescribeTableResult result = new DescribeTableResult().withTable(description);
    EasyMock.expect(dynamoDB.describeTable(tableName)).andReturn(result);
    final CreateTableRequest request = new CreateTableRequest().withAttributeDefinitions(ads).withKeySchema(kses)
        .withTableName(tableName);
    PowerMock.replayAll();
    assertEquals(description, DynamoDBManager.createTable(dynamoDB, request));
    PowerMock.verifyAll();
}
 

@Test(expected = IllegalStateException.class)
public void testCreateTableTableAlreadyExistsIncorrectKeySchema() {
    final Collection<AttributeDefinition> ads = Arrays.asList(new AttributeDefinition("Hash", ScalarAttributeType.S));
    final Collection<KeySchemaElement> kses = Arrays.asList(new KeySchemaElement("Hash", KeyType.HASH));
    final TableDescription description = new TableDescription().withAttributeDefinitions(ads).withKeySchema(kses)
        .withTableName(tableName);
    final DescribeTableResult result = new DescribeTableResult().withTable(description);
    EasyMock.expect(dynamoDB.describeTable(tableName)).andReturn(result);
    final Collection<AttributeDefinition> ads2 = Arrays.asList(new AttributeDefinition("Hash2", ScalarAttributeType.S));
    final Collection<KeySchemaElement> kses2 = Arrays.asList(new KeySchemaElement("Hash2", KeyType.HASH));
    final CreateTableRequest request = new CreateTableRequest().withAttributeDefinitions(ads2).withKeySchema(kses2)
        .withTableName(tableName);
    PowerMock.replayAll();
    DynamoDBManager.createTable(dynamoDB, request);
}
 

@Test
public void testGetTableStatus() {
    final TableDescription description = new TableDescription();
    final DescribeTableResult result = new DescribeTableResult().withTable(description);
    EasyMock.expect(dynamoDB.describeTable(tableName)).andReturn(result).anyTimes();
    for (final TableStatus status : TableStatus.values()) {
        description.setTableStatus(status);
        PowerMock.replayAll();
        assertEquals(status, DynamoDBManager.getTableStatus(dynamoDB, tableName));
        PowerMock.verifyAll();

    }
}
 

@Test
public void testTableExists() {
    final DescribeTableResult result = new DescribeTableResult();
    dynamoDB.describeTable(tableName);
    PowerMock.expectLastCall().andReturn(result);
    PowerMock.replayAll();
    assertTrue(DynamoDBManager.doesTableExist(dynamoDB, tableName));
    PowerMock.verifyAll();
}
 

@Test
public void testWaitForTableToBecomeActiveAlreadyActive() {
    final TableDescription table = new TableDescription();
    final DescribeTableResult result = new DescribeTableResult().withTable(table);
    table.setTableStatus(TableStatus.ACTIVE);
    dynamoDB.describeTable(tableName);
    PowerMock.expectLastCall().andReturn(result);
    PowerMock.expectLastCall().andReturn(result);
    PowerMock.replayAll();
    DynamoDBManager.waitForTableToBecomeActive(dynamoDB, tableName);
}
 

@Test
public void testWaitForTableToBecomeActiveCreatingThenActive() {
    // Creating table
    final TableDescription table1 = new TableDescription();
    table1.setTableStatus(TableStatus.CREATING);
    final DescribeTableResult result1 = new DescribeTableResult().withTable(table1);
    // Active table
    final TableDescription table2 = new TableDescription();
    table2.setTableStatus(TableStatus.ACTIVE);
    final DescribeTableResult result2 = new DescribeTableResult().withTable(table2);
    EasyMock.expect(dynamoDB.describeTable(tableName)).andReturn(result1);
    EasyMock.expect(dynamoDB.describeTable(tableName)).andReturn(result2);
    PowerMock.replayAll();
    DynamoDBManager.waitForTableToBecomeActive(dynamoDB, tableName);
}
 

@Test(expected = IllegalStateException.class)
public void testWaitForTableToBecomeActiveDeleting() {
    final TableDescription table = new TableDescription().withTableStatus(TableStatus.DELETING);
    final DescribeTableResult result = new DescribeTableResult().withTable(table);
    PowerMock.expectLastCall().andReturn(result);
    PowerMock.replayAll();
    DynamoDBManager.waitForTableToBecomeActive(dynamoDB, tableName);
}
 

@Test(expected = IllegalStateException.class)
public void testWaitForTableToBecomeActiveNeverGoingActive() {
    final TableDescription table = new TableDescription();
    final DescribeTableResult result = new DescribeTableResult().withTable(table);
    table.setTableStatus(TableStatus.CREATING);
    EasyMock.expect(dynamoDB.describeTable(tableName)).andReturn(result).anyTimes();
    PowerMock.replayAll();
    DynamoDBManager.waitForTableToBecomeActive(dynamoDB, tableName);
}
 

@Test
public void testWaitForTableToBecomeActiveUpdatingThenActive() {
    // Updating table
    final TableDescription table1 = new TableDescription();
    table1.setTableStatus(TableStatus.UPDATING);
    final DescribeTableResult result1 = new DescribeTableResult().withTable(table1);
    // Active table
    final TableDescription table2 = new TableDescription();
    table2.setTableStatus(TableStatus.ACTIVE);
    final DescribeTableResult result2 = new DescribeTableResult().withTable(table2);
    EasyMock.expect(dynamoDB.describeTable(tableName)).andReturn(result1);
    EasyMock.expect(dynamoDB.describeTable(tableName)).andReturn(result2);
    PowerMock.replayAll();
    DynamoDBManager.waitForTableToBecomeActive(dynamoDB, tableName);
}
 

@Test
public void getNoSqlTableMetaDataOk() {
    TableDescription tableDescription = new TableDescription().withTableArn(ARN).withTableStatus(ACTIVE_STATUS);
    DescribeTableResult describeResult = new DescribeTableResult().withTable(tableDescription);
    when(dynamoDb.describeTable(argThat((ArgumentMatcher<String>) argument -> true))).thenReturn(describeResult);
    NoSqlTableMetadataResponse result = underTest.getNoSqlTableMetaData(new NoSqlTableMetadataRequest());
    assertEquals(ARN, result.getId());
    assertEquals(ACTIVE_STATUS, result.getTableStatus());
    assertEquals(ResponseStatus.OK, result.getStatus());
}
 

protected boolean doesTableExist(String tableName) {
    try {
        DescribeTableRequest request = new DescribeTableRequest().withTableName(tableName);
        DescribeTableResult result = ddb.describeTable(request);
        return (result != null && "ACTIVE".equals(result.getTable().getTableStatus()));
    } catch (ResourceNotFoundException e) {
        return false;
    }
}
 

protected boolean doesTableExist(String tableName) {
    try {
        DescribeTableRequest request = new DescribeTableRequest().withTableName(tableName);
        DescribeTableResult result = ddb.describeTable(request);
        return (result != null && "ACTIVE".equals(result.getTable().getTableStatus()));
    } catch (ResourceNotFoundException e) {
        return false;
    }
}