package com.atlassian.messagequeue.internal.sqs;
   
   import com.amazonaws.services.sqs.model.SendMessageRequest;
   import com.atlassian.messagequeue.Message;
   import com.atlassian.messagequeue.MessageRunnerKey;
   import com.atlassian.messagequeue.internal.core.NestedMessage;
   import com.atlassian.messagequeue.registry.MessageContext;
   import org.junit.Test;
   import org.junit.runner.RunWith;
  import org.mockito.runners.MockitoJUnitRunner;
  
  import javax.annotation.Nullable;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  import java.util.Set;
  
  import static org.hamcrest.core.Is.is;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertThat;
  import static org.mockito.Matchers.any;
  import static org.mockito.Mockito.after;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  
  @RunWith(MockitoJUnitRunner.class)
  public class MultipleQueuesSQSMessageRunnerServiceTest extends AbstractSQSMessageRunnerServiceTest {
      private static final String QUEUE_NAME_KANGAROO = "kangarooQueue";
      private static final String QUEUE_NAME_KOALA = "koalaQueue";
      private static final MessageRunnerKey MESSAGE_RUNNER_KEY_MILK = MessageRunnerKey.of("milk");
      private static final MessageRunnerKey MESSAGE_RUNNER_KEY_BAMBOO = MessageRunnerKey.of("bamboo");
  
      private String queueUrlKangaroo;
      private String queueUrlKoala;
  
      @Override
      public void setUp() throws Exception {
          super.setUp();
  
          // Create 2 test queues
          sqsClient.createQueue(QUEUE_NAME_KANGAROO);
          queueUrlKangaroo = sqsClient.getQueueUrl(QUEUE_NAME_KANGAROO).getQueueUrl();
          sqsClient.createQueue(QUEUE_NAME_KOALA);
          queueUrlKoala = sqsClient.getQueueUrl(QUEUE_NAME_KOALA).getQueueUrl();
  
          // Register 2 message runners
          registryService.registerMessageRunner(MESSAGE_RUNNER_KEY_MILK, messageRunner);
          registryService.registerMessageRunner(MESSAGE_RUNNER_KEY_BAMBOO, messageRunner);
      }
  
      @Test
      public void messagesAreSentToCorrectQueues() {
          // Producer: milk -> kangaroo, else -> koala
          // Consumer: nothing
          messageRunnerService = withQueueConfigs(defaultMessageRunnerKeyToProducerMapper(), Collections.emptySet());
  
          messageRunnerService.addMessage(Message.create(MESSAGE_RUNNER_KEY_MILK, "m"));
          messageRunnerService.addMessage(Message.create(MESSAGE_RUNNER_KEY_BAMBOO, "b"));
          messageRunnerService.addMessage(Message.create(MESSAGE_RUNNER_KEY, "t"));
  
          verify(spyingSqsClient, times(3)).sendMessage(sendMessageRequestArgumentCaptor.capture());
          List<SendMessageRequest> requests = sendMessageRequestArgumentCaptor.getAllValues();
          // kangaroo queue
          NestedMessage nestedMessage = nestedMessageSerializer.deserialize(requests.get(0).getMessageBody());
          assertThat(requests.get(0).getQueueUrl(), is(queueUrlKangaroo));
          assertThat(nestedMessage.getPayload(), is("m"));
  
          // koala queue
          nestedMessage = nestedMessageSerializer.deserialize(requests.get(1).getMessageBody());
          assertThat(requests.get(1).getQueueUrl(), is(queueUrlKoala));
          assertThat(nestedMessage.getPayload(), is("b"));
  
          // default, koala queue
          nestedMessage = nestedMessageSerializer.deserialize(requests.get(2).getMessageBody());
          assertThat(requests.get(2).getQueueUrl(), is(queueUrlKoala));
          assertThat(nestedMessage.getPayload(), is("t"));
      }
  
      @Test(expected = NullPointerException.class)
      public void failImmediatelyIfNoProducerQueueConfig() {
          withQueueConfigs(null, Collections.emptySet());
      }
  
      @Test(expected = NullPointerException.class)
      public void failImmediatelyIfNoConsumerQueueConfig() {
          withQueueConfigs(getDefaultMessageRunnerKeyToProducerMapper(), null);
      }
  
      @Test
      public void onlyMatchedMessageIsConsumed() {
          // Listen to queue kangaroo only
          messageRunnerService = withQueueConfigs(defaultMessageRunnerKeyToProducerMapper(),
                  Collections.singleton(newConsumerQueueConfig(QUEUE_NAME_KANGAROO, queueUrlKangaroo)));
          messageRunnerService.initialiseMessageConsumers();
  
          // Target queue kangaroo, should be consumed
          messageRunnerService.addMessage(Message.create(MESSAGE_RUNNER_KEY_MILK, "m"));
          // Target default queue, should be ignored
          messageRunnerService.addMessage(Message.create(MESSAGE_RUNNER_KEY_BAMBOO, "b"));
 
         verify(spyingSqsClient, times(2)).sendMessage(sendMessageRequestArgumentCaptor.capture());
         List<SendMessageRequest> requests = sendMessageRequestArgumentCaptor.getAllValues();
         assertThat(requests.get(0).getQueueUrl(), is(queueUrlKangaroo));
         assertThat(requests.get(1).getQueueUrl(), is(queueUrlKoala));
         verify(messageRunner, after(1000).times(1)).processMessage(any(MessageContext.class));
         List<String> processedPayloads = new ArrayList<>();
         payloads.drainTo(processedPayloads);
         assertEquals(processedPayloads.size(), 1);
         assertThat(processedPayloads.get(0), is("m"));
     }
 
     private SQSMessageRunnerService withQueueConfigs(@Nullable SQSMessageRunnerKeyToProducerMapper messageRunnerKeyToProducerMapper,
                                                      @Nullable Set<SQSConsumerQueueConfig> consumerQueueConfigs) {
         return SQSMessageRunnerService.newBuilder(messageRunnerKeyToProducerMapper, consumerQueueConfigs)
                 .withAmazonSQSClient(spyingSqsClient)
                 .withReceiveWaitTimeSeconds(RECEIVE_WAIT_TIME_SECONDS)
                 .withMessageRunnerRegistryHelper(registryService)
                 .withTenantIdSetter(tenantIdSetter)
                 .withMessageInformationService(new SQSMessageInformationService(messageRunnerKeyToProducerMapper, tenantContextProvider, nestedMessageSerializer, tenantDataIdSupplier))
                 .withNestedMessageSerializer(nestedMessageSerializer)
                 .withSqsMessageVisibilityTimeoutManager(sqsMessageVisibilityTimeoutManager)
                 .withMessageValidatorRegistryHelper(messageValidatorRegistry).build();
 
     }
 
     private SQSConsumerQueueConfig newConsumerQueueConfig(String queueName, String queueUrl) {
         return new SQSConsumerQueueConfig() {
 
             @Override
             public String getQueueName() {
                 return queueName;
             }
 
             @Override
             public String getQueueUrl() {
                 return queueUrl;
             }
 
             @Override
             public int getCorePoolSize() {
                 return 2;
             }
 
             @Override
             public int getMaxPoolSize() {
                 return 2;
             }
 
             @Override
             public int getVisibilityExtensionPeriod() {
                 return 15;
             }
         };
     }
 
     private SQSMessageRunnerKeyToProducerMapper defaultMessageRunnerKeyToProducerMapper() {
         return messageRunnerKey -> {
             if (messageRunnerKey.equals(MESSAGE_RUNNER_KEY_MILK)) {
                 return new SQSProducerQueueConfig() {
                     @Override
                     public String getQueueName() {
                         return QUEUE_NAME_KANGAROO;
                     }
 
                     @Override
                     public String getQueueUrl() {
                         return queueUrlKangaroo;
                     }
                 };
             }
             return new SQSProducerQueueConfig() {
                 @Override
                 public String getQueueName() {
                     return QUEUE_NAME_KOALA;
                 }
 
                 @Override
                 public String getQueueUrl() {
                     return queueUrlKoala;
                 }
             };
         };
     }
 }