package com.jessecoyle;

import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
import com.amazonaws.services.dynamodbv2.model.PutItemRequest;
import com.amazonaws.services.dynamodbv2.model.PutItemResult;
import com.amazonaws.services.kms.AWSKMSClient;
import com.amazonaws.services.kms.model.GenerateDataKeyRequest;
import com.amazonaws.services.kms.model.GenerateDataKeyResult;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;
import org.mockito.internal.verification.VerificationModeFactory;

import java.nio.ByteBuffer;
import java.util.HashMap;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

public class JCredStashTest {

    private AmazonDynamoDBClient dynamoDBClient;
    private AWSKMSClient awskmsClient;

    @Before
    public void setUp()
    {
        dynamoDBClient = mock(AmazonDynamoDBClient.class);

        GenerateDataKeyResult generateDatakeyResult = new GenerateDataKeyResult();
        generateDatakeyResult.setCiphertextBlob(mock(ByteBuffer.class));
        generateDatakeyResult.setPlaintext(mock(ByteBuffer.class));

        awskmsClient = mock(AWSKMSClient.class);
        when(awskmsClient.generateDataKey(Mockito.any(GenerateDataKeyRequest.class))).thenReturn(generateDatakeyResult);
    }

    @Test
    public void testPutSecretDefaultVersion()
    {
        final PutItemRequest[] putItemRequest = new PutItemRequest[1];
        when(dynamoDBClient.putItem(Mockito.any(PutItemRequest.class))).thenAnswer(invocationOnMock -> {
            Object[] args = invocationOnMock.getArguments();
            putItemRequest[0] = (PutItemRequest) args[0];
            return new PutItemResult();
        });

        JCredStash credStash = new JCredStash(dynamoDBClient, awskmsClient);
        credStash.putSecret("table", "mysecret", "foo", "alias/foo", new HashMap<>(), null);

        verify(dynamoDBClient, VerificationModeFactory.times(1)).putItem(Mockito.any(PutItemRequest.class));
        Assert.assertEquals(putItemRequest[0].getItem().get("version").getS(), padVersion(1));
    }

    @Test
    public void testPutSecretNewVersion()
    {
        String version = "foover";
        final PutItemRequest[] putItemRequest = new PutItemRequest[1];
        when(dynamoDBClient.putItem(Mockito.any(PutItemRequest.class))).thenAnswer(invocationOnMock -> {
            Object[] args = invocationOnMock.getArguments();
            putItemRequest[0] = (PutItemRequest) args[0];
            return new PutItemResult();
        });

        JCredStash credStash = new JCredStash(dynamoDBClient, awskmsClient);
        credStash.putSecret("table", "mysecret", "foo", "alias/foo", new HashMap<>(), version);

        verify(dynamoDBClient, VerificationModeFactory.times(1)).putItem(Mockito.any(PutItemRequest.class));
        Assert.assertEquals(putItemRequest[0].getItem().get("version").getS(), version);
    }

    @Test
    public void testPutSecretAutoIncrementVersion()
    {
        final PutItemRequest[] putItemRequest = new PutItemRequest[1];
        when(dynamoDBClient.putItem(Mockito.any(PutItemRequest.class))).thenAnswer(invocationOnMock -> {
            Object[] args = invocationOnMock.getArguments();
            putItemRequest[0] = (PutItemRequest) args[0];
            return new PutItemResult();
        });

        JCredStash credStash = Mockito.spy(new JCredStash(dynamoDBClient, awskmsClient));
        Mockito.doReturn(padVersion(1)).when(credStash).getHighestVersion("table", "mysecret");
        credStash.putSecret("table", "mysecret", "foo", "alias/foo", new HashMap<>());

        verify(dynamoDBClient, VerificationModeFactory.times(1)).putItem(Mockito.any(PutItemRequest.class));
        Assert.assertEquals(putItemRequest[0].getItem().get("version").getS(), padVersion(2));
    }

    private String padVersion(int v)
    {
        return String.format("%019d", v);
    }
}

@@ -608,131 +608,8 @@
 p;%0A%0A
-import static org.mockito.Mockito.mock;%0Aimport static org.mockito.Mockito.verify;%0Aimport static org.mockito.Mockito.when;%0A%0A
 publ
@@ -787,24 +787,32 @@
 oDBClient = 
+Mockito.
 mock(AmazonD
@@ -965,16 +965,24 @@
 extBlob(
+Mockito.
 mock(Byt
@@ -1041,16 +1041,24 @@
 aintext(
+Mockito.
 mock(Byt
@@ -1098,16 +1098,24 @@
 lient = 
+Mockito.
 mock(AWS
@@ -1136,24 +1136,32 @@
 s);%0A        
+Mockito.
 when(awskmsC
@@ -1401,32 +1401,40 @@
 est%5B1%5D;%0A        
+Mockito.
 when(dynamoDBCli
@@ -1857,32 +1857,40 @@
 null);%0A%0A        
+Mockito.
 verify(dynamoDBC
@@ -2240,32 +2240,40 @@
 est%5B1%5D;%0A        
+Mockito.
 when(dynamoDBCli
@@ -2699,32 +2699,40 @@
 sion);%0A%0A        
+Mockito.
 verify(dynamoDBC
@@ -3059,16 +3059,24 @@
         
+Mockito.
 when(dyn
@@ -3618,16 +3618,24 @@
         
+Mockito.
 verify(d

Remove static imports in tests for consistency