package com.atlassian.messagequeue.internal.sqs;
   
   import static org.hamcrest.CoreMatchers.is;
   import static org.junit.Assert.assertThat;
   import static org.mockito.Mockito.*;
   
   import com.amazonaws.services.sqs.AmazonSQSClient;
   import com.amazonaws.services.sqs.model.MessageNotInflightException;
   import org.hamcrest.CoreMatchers;
  import org.hamcrest.Matchers;
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  import org.mockito.Mock;
  import org.junit.Rule;
  import org.mockito.junit.MockitoJUnit;
  import org.mockito.junit.MockitoRule;
  
  import java.util.concurrent.Future;
  
  public class SQSMessageVisibilityTimeoutManagerTest {
      @Rule
      public MockitoRule rule = MockitoJUnit.rule();
  
      @Mock
      AmazonSQSClient sqsClient;
      @Mock
      SQSConsumerQueueConfig queueConfig;
      private SQSMessageVisibilityTimeoutManager manager;
      private int baseVisibilityTimeoutSeconds;
  
      @Before
      public void setUp() throws Exception {
          baseVisibilityTimeoutSeconds = 1;
          manager = new SQSMessageVisibilityTimeoutManager(sqsClient, 1);
          when(queueConfig.getVisibilityExtensionPeriod()).thenReturn(1);
      }
  
      @Test
      public void stopExtendingVisibilityTimeoutOnException() throws Exception {
          doThrow(new MessageNotInflightException("message not in flight")).when(sqsClient).changeMessageVisibility(anyString(), anyString(), anyInt());
  
          final Future<?> future = manager.scheduleVisibilityTimeoutExtension(queueConfig, "receiptHandle");
  
          verify(sqsClient, after(baseVisibilityTimeoutSeconds * 1000 * 2).times(1)).changeMessageVisibility(anyString(), anyString(), anyInt());
          assertThat(future.isDone(), is(true));
          assertThat(manager.visibilityTimeoutExtensionWorkerPool.getQueue(), Matchers.empty());
      }
  
      @Test
      public void visibilityTimeoutExtensionRecurs() throws Exception {
          manager.scheduleVisibilityTimeoutExtension(queueConfig, "receiptHandle");
  
          // Within 2500ms, the extension task is enabled after a delay of 1000ms and then again at 2000ms producing 2 wanted invocations
          final int wantedNumberOfInvocations = 2;
          verify(sqsClient, after(2500).times(wantedNumberOfInvocations)).changeMessageVisibility(anyString(), anyString(), anyInt());
      }
  
      @After
      public void tearDown() throws Exception {
          manager.shutdown();
      }
  }