package org.sagebionetworks.kinesis;

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

import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Stream;

import com.amazonaws.services.kinesisfirehose.AmazonKinesisFirehoseClient;
import com.amazonaws.services.kinesisfirehose.model.PutRecordBatchRequest;
import com.amazonaws.services.kinesisfirehose.model.PutRecordRequest;
import com.amazonaws.services.kinesisfirehose.model.Record;
import com.google.common.collect.Lists;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;
import org.sagebionetworks.StackConfiguration;

@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.LENIENT)
class AwsKinesisFirehoseLoggerImplTest {

	@Mock
	StackConfiguration mockStackConfig;

	@Mock
	AmazonKinesisFirehoseClient mockKinesisFirehoseClient;

	@InjectMocks
	AwsKinesisFirehoseLoggerImpl kinesisFirehoseLogger;

	@Mock
	AwsKinesisLogRecord mockRecord1;

	@Mock
	AwsKinesisLogRecord mockRecord2;


	String kinesisStreamSuffix = "myKinesisStream";
	String stack = "dev";
	String instance = "test";


	@BeforeEach
	public void setup(){
		when(mockStackConfig.getStack()).thenReturn(stack);
		when(mockRecord1.withStack(stack)).thenReturn(mockRecord1);
		when(mockRecord2.withStack(stack)).thenReturn(mockRecord2);
		when(mockStackConfig.getStackInstance()).thenReturn(instance);
	}


	@Test
	public void testLogBatch(){
		List<AwsKinesisLogRecord> mockRecordList = Lists.newArrayList(mockRecord1, mockRecord2);


		byte[] mockRecord1Bytes = new byte[1];
		byte[] mockRecord2Bytes = new byte[2];
		when(mockRecord1.toBytes()).thenReturn(mockRecord1Bytes);
		when(mockRecord2.toBytes()).thenReturn(mockRecord2Bytes);


		//method under test
		kinesisFirehoseLogger.logBatch(kinesisStreamSuffix, mockRecordList);

		verify(mockRecord1).withStack(stack);
		verify(mockRecord1).withInstance(instance);
		verify(mockRecord2).withStack(stack);
		verify(mockRecord2).withInstance(instance);

		verify(mockKinesisFirehoseClient).putRecordBatch(
				new PutRecordBatchRequest()
						.withDeliveryStreamName("devtestmyKinesisStream")
						.withRecords(Arrays.asList(
								new Record().withData(ByteBuffer.wrap(mockRecord1Bytes)),
								new Record().withData(ByteBuffer.wrap(mockRecord2Bytes))))
		);

		verifyNoMoreInteractions(mockKinesisFirehoseClient);
	}

	@Test
	public void testPartitioning(){
		List<AwsKinesisLogRecord> repeatList = Collections.nCopies(502, mockRecord1);

		byte[] mockRecord1Bytes = new byte[1];
		when(mockRecord1.toBytes()).thenReturn(mockRecord1Bytes);

		//method under test
		kinesisFirehoseLogger.logBatch(kinesisStreamSuffix, repeatList);


		//verify that 2 calls were made
		verify(mockKinesisFirehoseClient).putRecordBatch(
				new PutRecordBatchRequest()
						.withDeliveryStreamName("devtestmyKinesisStream")
						.withRecords(Collections.nCopies(500,
								new Record().withData(ByteBuffer.wrap(mockRecord1Bytes))))
		);

		verify(mockKinesisFirehoseClient).putRecordBatch(
				new PutRecordBatchRequest()
						.withDeliveryStreamName("devtestmyKinesisStream")
						.withRecords(Collections.nCopies(2,
								new Record().withData(ByteBuffer.wrap(mockRecord1Bytes))))
		);

		verifyNoMoreInteractions(mockKinesisFirehoseClient);
	}
}
@@ -61,20 +61,28 @@
 ito.
-Mockito.mock
+ArgumentMatchers.any
 ;%0Aim
@@ -332,360 +332,8 @@
 st;%0A
-import java.util.stream.Stream;%0A%0Aimport com.amazonaws.services.kinesisfirehose.AmazonKinesisFirehoseClient;%0Aimport com.amazonaws.services.kinesisfirehose.model.PutRecordBatchRequest;%0Aimport com.amazonaws.services.kinesisfirehose.model.PutRecordRequest;%0Aimport com.amazonaws.services.kinesisfirehose.model.Record;%0Aimport com.google.common.collect.Lists;
 %0Aimp
@@ -701,16 +701,341 @@
 ation;%0A%0A
+import com.amazonaws.services.kinesisfirehose.AmazonKinesisFirehoseClient;%0Aimport com.amazonaws.services.kinesisfirehose.model.PutRecordBatchRequest;%0Aimport com.amazonaws.services.kinesisfirehose.model.PutRecordBatchResult;%0Aimport com.amazonaws.services.kinesisfirehose.model.Record;%0Aimport com.google.common.collect.Lists;%0A%0A
 @ExtendW
@@ -1414,16 +1414,65 @@
 ecord2;%0A
+%09%0A%09@Mock%0A%09PutRecordBatchResult mockRecordResult;%0A
 %0A%0A%09Strin
@@ -1844,16 +1844,102 @@
 tance);%0A
+%09%09when(mockKinesisFirehoseClient.putRecordBatch(any())).thenReturn(mockRecordResult);%0A
 %09%7D%0A%0A%0A%09@T

Fix failing test