Examples with Shard com.amazonaws.services.kinesis.model.Shard used on opensource projects

Example 11 with Shard

use of com.amazonaws.services.kinesis.model.Shard in project flink by apache.

the class FlinkKinesisConsumerTest method testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint.

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint() throws Exception {
    KinesisDataFetcher mockedFetcher = Mockito.mock(KinesisDataFetcher.class);
    PowerMockito.whenNew(KinesisDataFetcher.class).withAnyArguments().thenReturn(mockedFetcher);
    // assume the given config is correct
    PowerMockito.mockStatic(KinesisConfigUtil.class);
    PowerMockito.doNothing().when(KinesisConfigUtil.class);
    HashMap<KinesisStreamShard, SequenceNumber> fakeRestoredState = new HashMap<>();
    fakeRestoredState.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))), new SequenceNumber(UUID.randomUUID().toString()));
    fakeRestoredState.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))), new SequenceNumber(UUID.randomUUID().toString()));
    fakeRestoredState.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))), new SequenceNumber(UUID.randomUUID().toString()));
    fakeRestoredState.put(new KinesisStreamShard("fakeStream2", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))), new SequenceNumber(UUID.randomUUID().toString()));
    fakeRestoredState.put(new KinesisStreamShard("fakeStream2", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))), new SequenceNumber(UUID.randomUUID().toString()));
    TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
    consumer.restoreState(fakeRestoredState);
    consumer.open(new Configuration());
    consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
    Mockito.verify(mockedFetcher).setIsRestoringFromFailure(true);
    for (Map.Entry<KinesisStreamShard, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
        Mockito.verify(mockedFetcher).advanceLastDiscoveredShardOfStream(restoredShard.getKey().getStreamName(), restoredShard.getKey().getShard().getShardId());
        Mockito.verify(mockedFetcher).registerNewSubscribedShardState(new KinesisStreamShardState(restoredShard.getKey(), restoredShard.getValue()));
    }
}

Example 12 with Shard

use of com.amazonaws.services.kinesis.model.Shard in project flink by apache.

the class KinesisDataFetcherTest method testStreamToLastSeenShardStateIsCorrectlySetWhenNewShardsFoundSinceRestoredCheckpoint.

@Test
public void testStreamToLastSeenShardStateIsCorrectlySetWhenNewShardsFoundSinceRestoredCheckpoint() throws Exception {
    List<String> fakeStreams = new LinkedList<>();
    fakeStreams.add("fakeStream1");
    fakeStreams.add("fakeStream2");
    Map<KinesisStreamShard, String> restoredStateUnderTest = new HashMap<>();
    // fakeStream1 has 3 shards before restore
    restoredStateUnderTest.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))), UUID.randomUUID().toString());
    restoredStateUnderTest.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))), UUID.randomUUID().toString());
    restoredStateUnderTest.put(new KinesisStreamShard("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))), UUID.randomUUID().toString());
    // fakeStream2 has 2 shards before restore
    restoredStateUnderTest.put(new KinesisStreamShard("fakeStream2", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0))), UUID.randomUUID().toString());
    restoredStateUnderTest.put(new KinesisStreamShard("fakeStream2", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1))), UUID.randomUUID().toString());
    Map<String, Integer> streamToShardCount = new HashMap<>();
    // fakeStream1 had 3 shards before & 1 new shard after restore
    streamToShardCount.put("fakeStream1", 3 + 1);
    // fakeStream2 had 2 shards before & 3 new shard after restore
    streamToShardCount.put("fakeStream2", 2 + 3);
    HashMap<String, String> subscribedStreamsToLastSeenShardIdsUnderTest = KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(fakeStreams);
    // using a non-resharded streams kinesis behaviour to represent that Kinesis is not resharded AFTER the restore
    final TestableKinesisDataFetcher fetcher = new TestableKinesisDataFetcher(fakeStreams, new Properties(), 10, 2, new AtomicReference<Throwable>(), new LinkedList<KinesisStreamShardState>(), subscribedStreamsToLastSeenShardIdsUnderTest, FakeKinesisBehavioursFactory.nonReshardedStreamsBehaviour(streamToShardCount));
    for (Map.Entry<KinesisStreamShard, String> restoredState : restoredStateUnderTest.entrySet()) {
        fetcher.advanceLastDiscoveredShardOfStream(restoredState.getKey().getStreamName(), restoredState.getKey().getShard().getShardId());
        fetcher.registerNewSubscribedShardState(new KinesisStreamShardState(restoredState.getKey(), new SequenceNumber(restoredState.getValue())));
    }
    fetcher.setIsRestoringFromFailure(true);
    PowerMockito.whenNew(ShardConsumer.class).withAnyArguments().thenReturn(Mockito.mock(ShardConsumer.class));
    Thread runFetcherThread = new Thread(new Runnable() {

        @Override
        public void run() {
            try {
                fetcher.runFetcher();
            } catch (Exception e) {
            //
            }
        }
    });
    runFetcherThread.start();
    // sleep a while before closing
    Thread.sleep(1000);
    fetcher.shutdownFetcher();
    // assert that the streams tracked in the state are identical to the subscribed streams
    Set<String> streamsInState = subscribedStreamsToLastSeenShardIdsUnderTest.keySet();
    assertTrue(streamsInState.size() == fakeStreams.size());
    assertTrue(streamsInState.containsAll(fakeStreams));
    // assert that the last seen shards in state is correctly set
    for (Map.Entry<String, String> streamToLastSeenShard : subscribedStreamsToLastSeenShardIdsUnderTest.entrySet()) {
        assertTrue(streamToLastSeenShard.getValue().equals(KinesisShardIdGenerator.generateFromShardOrder(streamToShardCount.get(streamToLastSeenShard.getKey()) - 1)));
    }
}

Example 13 with Shard

use of com.amazonaws.services.kinesis.model.Shard in project flink by apache.

the class ShardConsumerTest method testCorrectNumOfCollectedRecordsAndUpdatedState.

@Test
public void testCorrectNumOfCollectedRecordsAndUpdatedState() {
    KinesisStreamShard fakeToBeConsumedShard = new KinesisStreamShard("fakeStream", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)).withHashKeyRange(new HashKeyRange().withStartingHashKey("0").withEndingHashKey(new BigInteger(StringUtils.repeat("FF", 16), 16).toString())));
    LinkedList<KinesisStreamShardState> subscribedShardsStateUnderTest = new LinkedList<>();
    subscribedShardsStateUnderTest.add(new KinesisStreamShardState(fakeToBeConsumedShard, new SequenceNumber("fakeStartingState")));
    TestableKinesisDataFetcher fetcher = new TestableKinesisDataFetcher(Collections.singletonList("fakeStream"), new Properties(), 10, 2, new AtomicReference<Throwable>(), subscribedShardsStateUnderTest, KinesisDataFetcher.createInitialSubscribedStreamsToLastDiscoveredShardsState(Collections.singletonList("fakeStream")), Mockito.mock(KinesisProxyInterface.class));
    new ShardConsumer<>(fetcher, 0, subscribedShardsStateUnderTest.get(0).getKinesisStreamShard(), subscribedShardsStateUnderTest.get(0).getLastProcessedSequenceNum(), FakeKinesisBehavioursFactory.totalNumOfRecordsAfterNumOfGetRecordsCalls(1000, 9)).run();
    assertTrue(fetcher.getNumOfElementsCollected() == 1000);
    assertTrue(subscribedShardsStateUnderTest.get(0).getLastProcessedSequenceNum().equals(SentinelSequenceNumber.SENTINEL_SHARD_ENDING_SEQUENCE_NUM.get()));
}

Example 14 with Shard

use of com.amazonaws.services.kinesis.model.Shard in project camel by apache.

the class KinesisConsumerTest method setup.

@Before
public void setup() throws Exception {
    KinesisEndpoint endpoint = new KinesisEndpoint(null, "streamName", component);
    endpoint.setAmazonKinesisClient(kinesisClient);
    endpoint.setIteratorType(ShardIteratorType.LATEST);
    undertest = new KinesisConsumer(endpoint, processor);
    when(kinesisClient.getRecords(any(GetRecordsRequest.class))).thenReturn(new GetRecordsResult().withNextShardIterator("nextShardIterator"));
    when(kinesisClient.describeStream(any(DescribeStreamRequest.class))).thenReturn(new DescribeStreamResult().withStreamDescription(new StreamDescription().withShards(new Shard().withShardId("shardId"))));
    when(kinesisClient.getShardIterator(any(GetShardIteratorRequest.class))).thenReturn(new GetShardIteratorResult().withShardIterator("shardIterator"));
}

Example 15 with Shard

use of com.amazonaws.services.kinesis.model.Shard in project beam by apache.

the class AmazonKinesisMock method describeStream.

@Override
public DescribeStreamResult describeStream(String streamName, String exclusiveStartShardId) {
    int nextShardId = 0;
    if (exclusiveStartShardId != null) {
        nextShardId = parseInt(exclusiveStartShardId) + 1;
    }
    boolean hasMoreShards = nextShardId + 1 < shardedData.size();
    List<Shard> shards = newArrayList();
    if (nextShardId < shardedData.size()) {
        shards.add(new Shard().withShardId(Integer.toString(nextShardId)));
    }
    return new DescribeStreamResult().withStreamDescription(new StreamDescription().withHasMoreShards(hasMoreShards).withShards(shards));
}