Examples with DefaultRateLimiter org.apache.beam.sdk.io.aws2.kinesis.RateLimitPolicyFactory.DefaultRateLimiter used on opensource projects

Example 1 with DefaultRateLimiter

use of org.apache.beam.sdk.io.aws2.kinesis.RateLimitPolicyFactory.DefaultRateLimiter in project beam by apache.

the class RateLimitPolicyFactoryTest method defaultRateLimiterShouldUseBackoffs.

@Test
public void defaultRateLimiterShouldUseBackoffs() throws Exception {
    assertThat(withDefaultRateLimiter().getRateLimitPolicy()).isInstanceOf(DefaultRateLimiter.class);
    assertThat(withDefaultRateLimiter(millis(1), millis(1), millis(1)).getRateLimitPolicy()).isInstanceOf(DefaultRateLimiter.class);
    Sleeper sleeper = mock(Sleeper.class);
    BackOff emptySuccess = mock(BackOff.class);
    BackOff throttled = mock(BackOff.class);
    RateLimitPolicy policy = new DefaultRateLimiter(emptySuccess, throttled, sleeper);
    // reset emptySuccess after receiving at least 1 record, throttled is reset on any success
    policy.onSuccess(ImmutableList.of(mock(KinesisRecord.class)));
    verify(emptySuccess).reset();
    verify(throttled).reset();
    verifyNoInteractions(sleeper);
    clearInvocations(emptySuccess, throttled);
    when(emptySuccess.nextBackOffMillis()).thenReturn(88L, 99L);
    // throttle if no records received, throttled is reset again
    policy.onSuccess(ImmutableList.of());
    policy.onSuccess(ImmutableList.of());
    verify(emptySuccess, times(2)).nextBackOffMillis();
    verify(throttled, times(2)).reset();
    verify(sleeper).sleep(88L);
    verify(sleeper).sleep(99L);
    verifyNoMoreInteractions(sleeper, throttled, emptySuccess);
    clearInvocations(emptySuccess, throttled, sleeper);
    when(throttled.nextBackOffMillis()).thenReturn(111L, 222L);
    // throttle onThrottle
    policy.onThrottle(mock(KinesisClientThrottledException.class));
    policy.onThrottle(mock(KinesisClientThrottledException.class));
    verify(throttled, times(2)).nextBackOffMillis();
    verify(sleeper).sleep(111L);
    verify(sleeper).sleep(222L);
    verifyNoMoreInteractions(sleeper, throttled, emptySuccess);
}