Java Code Examples for org.mockito.invocation.InvocationOnMock#getArgumentAt()

/**
 * Test set up.
 */
@Before
public void setUp() {
    @SuppressWarnings("unchecked")
    final SessionIdentifierFactory<String> factory = Mockito
            .mock(SessionIdentifierFactory.class);
    final Answer<String> answer = new Answer<String>() {
        @Override
        public String answer(InvocationOnMock invocation) throws Throwable {
            return invocation.getArgumentAt(0, String.class) + "-G";
        }
    };
    Mockito.when(factory.createSessionIdentifier(Mockito.anyString()))
            .then(answer);
    sessionId = factory.createSessionIdentifier("dummy");
    storage = new SessionStorage<String>(factory);
}
 

private void testStartMicroSpilling(boolean useSplaySort) throws Exception{
  final ExternalSortTracer tracer = new ExternalSortTracer();
  try (MemoryRun memoryRun = new MemoryRun(externalSort, producer, allocator, generator.getSchema(), tracer,
    2, useSplaySort, 8192, mock(ExecutionControls.class))) {
    final int totalAdded = addBatches(memoryRun);
    final DiskRunManager diskRunManager = mock(DiskRunManager.class);

    //the hypercontainer should be sorted before micro-spilling starts
    Answer<Void> spillAnswer = new Answer<Void>() {
      @Override
      public Void answer(InvocationOnMock invocation) throws Throwable {
        VectorContainer hyperBatch = invocation.getArgumentAt(0, VectorContainer.class);
        validateHyperBatch(hyperBatch, totalAdded);
        return null;
      }
    };

    Mockito.doAnswer(spillAnswer).when(diskRunManager).startMicroSpilling(Mockito.any(VectorContainer.class));

    memoryRun.startMicroSpilling(diskRunManager);
  }
}
 

@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
	SendMessageCommandV6 sendMessageCmd = invocation.getArgumentAt(1, SendMessageCommandV6.class);
	CommandProcessor commandProcessor = PlexusComponentLocator.lookup(CommandProcessor.class,
	      CommandType.RESULT_MESSAGE_SEND.toString());

	SendMessageResultCommandV6 resultCmd = new SendMessageResultCommandV6(sendMessageCmd.getMessageCount());
	resultCmd.correlate(sendMessageCmd);

	Map<Integer, SendMessageResult> results = new HashMap<>();
	for (List<ProducerMessage<?>> pmsgList : sendMessageCmd.getProducerMessages()) {
		for (ProducerMessage<?> pmsg : pmsgList) {
			results.put(pmsg.getMsgSeqNo(), new SendMessageResult(false, false, null));
		}
	}
	resultCmd.addResults(results);

	commandProcessor.process(new CommandProcessorContext(resultCmd, null));

	return null;
}
 

public AddressSelectorWrapper verifyNumberOfForceCachRefresh(int expectedNumber) {
    int count = 0;
    for (InvocationOnMock invocationOnMock : invocationOnMockList) {
        boolean forceRefresh;
        if (invocationOnMock.getMethod().getName().endsWith("resolveAllUriAsync")) {
            forceRefresh = invocationOnMock.getArgumentAt(2, Boolean.class);
        } else {
            forceRefresh = invocationOnMock.getArgumentAt(1, Boolean.class);
        }
        if (forceRefresh) {
            count++;
        }
    }
    assertThat(count).isEqualTo(expectedNumber);
    return this;
}
 

Answer getRegisteredTableAnswer() {
  return (InvocationOnMock x) -> {
    TableDescriptor descriptor = x.getArgumentAt(0, TableDescriptor.class);
    String tableId = "t1";
    TableImpl mockTable = mock(TableImpl.class);
    when(mockTable.getTableId()).thenReturn(tableId);
    this.registeredTables.putIfAbsent(descriptor.getTableId(), mockTable);
    return this.registeredTables.get(descriptor.getTableId());
  };
}
 

@Override
public Void answer(InvocationOnMock invocation)
    throws InvocationTargetException, InterruptedException {
  IRunnableWithProgress runnable = invocation.getArgumentAt(2, IRunnableWithProgress.class);
  runnable.run(new NullProgressMonitor());
  return null;
}
 

Answer getRegisterMessageStreamAnswer() {
  return (InvocationOnMock x) -> {
    Integer id = x.getArgumentAt(0, Integer.class);
    MessageStream stream = x.getArgumentAt(1, MessageStream.class);
    registeredStreams.put(id, stream);
    return null;
  };
}
 

@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
    String name = invocation.getArgumentAt(0, String.class);
    String value = invocation.getArgumentAt(1, String.class);
    apply(name, value);
    return null;
}
 

private static <T> Answer<T> returnConverterValue() {
	return new Answer<T>() {
		@SuppressWarnings("unchecked")
		@Override
		public T answer(InvocationOnMock invocation) throws Throwable {
			return (T)invocation.getArgumentAt(0, Object.class);
		}
	};
}
 

private GetMetricDataResult mockMetricData(InvocationOnMock invocation, int numMetrics, int numSamples)
{
    GetMetricDataRequest request = invocation.getArgumentAt(0, GetMetricDataRequest.class);

    /**
     * Confirm that all available criteria were pushed down into Cloudwatch Metrics
     */
    List<MetricDataQuery> queries = request.getMetricDataQueries();
    assertEquals(1, queries.size());
    MetricDataQuery query = queries.get(0);
    MetricStat stat = query.getMetricStat();
    assertEquals("m1", query.getId());
    assertNotNull(stat.getPeriod());
    assertNotNull(stat.getMetric());
    assertNotNull(stat.getStat());
    assertNotNull(stat.getMetric().getMetricName());
    assertNotNull(stat.getMetric().getNamespace());
    assertNotNull(stat.getMetric().getDimensions());
    assertEquals(1, stat.getMetric().getDimensions().size());

    String nextToken = (request.getNextToken() == null) ? "valid" : null;
    List<MetricDataResult> samples = new ArrayList<>();

    for (int i = 0; i < numMetrics; i++) {
        List<Double> values = new ArrayList<>();
        List<Date> timestamps = new ArrayList<>();
        for (double j = 0; j < numSamples; j++) {
            values.add(j);
            timestamps.add(new Date(System.currentTimeMillis() + (int) j));
        }
        samples.add(new MetricDataResult().withValues(values).withTimestamps(timestamps).withId("m1"));
    }

    return new GetMetricDataResult().withNextToken(nextToken).withMetricDataResults(samples);
}
 

@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
	SendMessageCommandV6 sendMessageCmd = invocation.getArgumentAt(1, SendMessageCommandV6.class);
	CommandProcessor commandProcessor = PlexusComponentLocator.lookup(CommandProcessor.class,
	      CommandType.ACK_MESSAGE_SEND.toString());

	SendMessageAckCommandV6 acceptCmd = new SendMessageAckCommandV6();
	acceptCmd.setSuccess(true);
	acceptCmd.correlate(sendMessageCmd);

	commandProcessor.process(new CommandProcessorContext(acceptCmd, null));

	return null;
}
 

@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
	SendMessageCommandV6 sendMessageCmd = invocation.getArgumentAt(1, SendMessageCommandV6.class);
	CommandProcessor commandProcessor = PlexusComponentLocator.lookup(CommandProcessor.class,
	      CommandType.ACK_MESSAGE_SEND.toString());

	SendMessageAckCommandV6 acceptCmd = new SendMessageAckCommandV6();
	acceptCmd.setSuccess(false);
	acceptCmd.correlate(sendMessageCmd);

	commandProcessor.process(new CommandProcessorContext(acceptCmd, null));

	return null;
}
 

protected Object assertConsumeTextMessage(CountDownLatch latch, String expectedMessage, InvocationOnMock invocation) {
    final WebSocketSessionInfo sessionInfo = invocation.getArgumentAt(0, WebSocketSessionInfo.class);
    assertNotNull(sessionInfo.getLocalAddress());
    assertNotNull(sessionInfo.getRemoteAddress());
    assertNotNull(sessionInfo.getSessionId());
    assertEquals(isSecure(), sessionInfo.isSecure());

    final String receivedMessage = invocation.getArgumentAt(1, String.class);
    assertNotNull(receivedMessage);
    assertEquals(expectedMessage, receivedMessage);
    latch.countDown();
    return null;
}
 

static void capture(List<Pair<Uri, RxDocumentServiceRequest>> capturedRequests, InvocationOnMock invocation) {
     Uri physicalUri = invocation.getArgumentAt(0, Uri.class);
     RxDocumentServiceRequest request = invocation.getArgumentAt(1, RxDocumentServiceRequest.class);
     logger.debug("Uri: {}, request {}", physicalUri, request);
     capturedRequests.add(Pair.of(physicalUri, request));
}
 

/**
 * Test method to activate a grammar with a field scope and a form scope.
 * 
 * @throws Exception
 *             Test failed.
 * @throws JVoiceXMLEvent
 *             Test failed.
 */
@Test(timeout = 5000)
public void testActivateFormGrammarsInField() throws Exception,
        JVoiceXMLEvent {
    final VoiceXmlDocument doc = new VoiceXmlDocument();
    final Vxml vxml = doc.getVxml();
    final Form form = vxml.appendChild(Form.class);
    final Grammar grammar = form.appendChild(Grammar.class);
    grammar.setVersion("1.0");
    grammar.setType(GrammarType.SRGS_XML);
    final Rule rule = grammar.appendChild(Rule.class);
    final OneOf oneof = rule.appendChild(OneOf.class);
    final Item item1 = oneof.appendChild(Item.class);
    item1.addText("visa");
    final Item item2 = oneof.appendChild(Item.class);
    item2.addText("mastercard");
    final Item item3 = oneof.appendChild(Item.class);
    item3.addText("american express");
    form.appendChild(Field.class);
    final Dialog executableForm = new ExecutablePlainForm();
    executableForm.setNode(form);
    final FormInterpretationAlgorithm fia = new FormInterpretationAlgorithm(
            context, interpreter, executableForm);
    final Object lock = new Object();
    final Thread thread = new Thread() {
        @Override
        public void run() {
            try {
                fia.initialize(profile, null);
                fia.mainLoop();
            } catch (JVoiceXMLEvent e) {
                Assert.fail(e.getMessage());
            }
        };
    };
    thread.start();
    final UserInput input = platform.getUserInput();
    final Answer<Integer> answer = new Answer<Integer>() {
        @Override
        public Integer answer(final InvocationOnMock invocation)
                throws Throwable {
            final Collection grammars = invocation.getArgumentAt(0,
                    Collection.class);
            synchronized (lock) {
                lock.notifyAll();
            }
            return grammars.size();
        }
    };
    Mockito.when(
            input.activateGrammars(Mockito
                    .anyCollectionOf(GrammarDocument.class))).then(answer);
    synchronized (lock) {
        lock.wait();
    }
    Mockito.verify(input).activateGrammars(
            Mockito.anyCollectionOf(GrammarDocument.class));
}
 

@Test
public void preemptDescendantsOnly() {

  final int MAX_TASKS = 2;
  TezConfiguration tezConf = new TezConfiguration();
  tezConf.setInt(TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS, MAX_TASKS);

  ApplicationId appId = ApplicationId.newInstance(2000, 1);
  ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
  Long parentTask1 = new Long(1);
  Long parentTask2 = new Long(2);
  Long childTask1 = new Long(3);
  Long grandchildTask1 = new Long(4);

  TaskSchedulerContext
      mockContext = TestTaskSchedulerHelpers.setupMockTaskSchedulerContext("", 0, "", true,
      appAttemptId, 1000l, null, tezConf);
  when(mockContext.getVertexIndexForTask(parentTask1)).thenReturn(0);
  when(mockContext.getVertexIndexForTask(parentTask2)).thenReturn(0);
  when(mockContext.getVertexIndexForTask(childTask1)).thenReturn(1);
  when(mockContext.getVertexIndexForTask(grandchildTask1)).thenReturn(2);

  DagInfo mockDagInfo = mock(DagInfo.class);
  when(mockDagInfo.getTotalVertices()).thenReturn(3);
  BitSet vertex1Descendants = new BitSet();
  vertex1Descendants.set(1);
  vertex1Descendants.set(2);
  BitSet vertex2Descendants = new BitSet();
  vertex2Descendants.set(2);
  BitSet vertex3Descendants = new BitSet();
  when(mockDagInfo.getVertexDescendants(0)).thenReturn(vertex1Descendants);
  when(mockDagInfo.getVertexDescendants(1)).thenReturn(vertex2Descendants);
  when(mockDagInfo.getVertexDescendants(2)).thenReturn(vertex3Descendants);
  when(mockContext.getCurrentDagInfo()).thenReturn(mockDagInfo);

  Priority priority1 = Priority.newInstance(1);
  Priority priority2 = Priority.newInstance(2);
  Priority priority3 = Priority.newInstance(3);
  Priority priority4 = Priority.newInstance(4);
  Resource resource = Resource.newInstance(1024, 1);

  MockLocalTaskSchedulerSerivce taskSchedulerService = new MockLocalTaskSchedulerSerivce(mockContext);

  // The mock context need to send a deallocate container request to the scheduler service
  Answer<Void> answer = new Answer<Void>() {
    @Override
    public Void answer(InvocationOnMock invocation) {
      ContainerId containerId = invocation.getArgumentAt(0, ContainerId.class);
      taskSchedulerService.deallocateContainer(containerId);
      return null;
    }
  };
  doAnswer(answer).when(mockContext).preemptContainer(any(ContainerId.class));

  taskSchedulerService.initialize();
  taskSchedulerService.start();
  taskSchedulerService.startRequestHandlerThread();

  MockAsyncDelegateRequestHandler requestHandler = taskSchedulerService.getRequestHandler();
  taskSchedulerService.allocateTask(parentTask1, resource, null, null, priority1, null, null);
  taskSchedulerService.allocateTask(childTask1, resource, null, null, priority3, null, null);
  taskSchedulerService.allocateTask(grandchildTask1, resource, null, null, priority4, null, null);
  requestHandler.drainRequest(3);

  // We should not preempt if we have not reached max task allocations
  Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, requestHandler.allocateCount);
  Assert.assertTrue("Another allocation should not fit", !requestHandler.shouldProcess());

  // Next task allocation should preempt
  taskSchedulerService.allocateTask(parentTask2, Resource.newInstance(1024, 1), null, null, priority2, null, null);
  requestHandler.drainRequest(5);

  // All allocated tasks should have been removed
  Assert.assertEquals("Wrong number of preempted tasks", 1, requestHandler.preemptCount);
}
 

@Test
public void testCompressedBuffer() throws TransportException, IllegalStateException, IOException {
  /*
   * Create mock client, requets, and replies
   */
  AmazonS3Client mockClient = getMockClient();

  /*
   * Capture the InputStream into a ByteArrayOutputStream before the Transport thread closes the
   * InputStream and makes it unavailable for reading.
   */
  ByteArrayOutputStream captured = new ByteArrayOutputStream();
  Answer answer = new Answer() {
    @Override
    public Object answer(InvocationOnMock invocation) throws Throwable {
      UploadPartRequest req = invocation.getArgumentAt(0, UploadPartRequest.class);
      captured.write(req.getInputStream());
      return new UploadPartResult();
    }
  };

  Mockito.doAnswer(answer).when(mockClient).uploadPart(any(UploadPartRequest.class));

  /*
   * Fill buffer with mock data
   */
  S3TransportBuffer buffer = new S3TransportBuffer(1000, true, new S3TransportSerializer());
  InternalEvent mockIevent = mock(InternalEvent.class);
  doReturn("foo").when(mockIevent).getSerialized();

  /*
   * Create transport
   */
  Map<String, MultiPartUpload> multiPartUploads = new HashMap<String, MultiPartUpload>(0);
  S3Transport transport =
      new S3Transport(mockClient, "bucket", "basepath", true, multiPartUploads);

  /*
   * Do actual test
   */
  buffer.add(mockIevent);
  LinkedHashMap<String, String> partitions = new LinkedHashMap<String, String>();
  partitions.put(S3Transport.FILENAME_KEY, "a_filename");
  ArgumentCaptor<UploadPartRequest> argument = ArgumentCaptor.forClass(UploadPartRequest.class);

  buffer.close();
  transport.sendBatch(buffer, partitions, new TestContext());
  verify(mockClient).uploadPart(argument.capture());

  /*
   * Check results
   */
  assertEquals("bucket", argument.getValue().getBucketName());
  assertEquals("basepath/a_filename.bz2", argument.getValue().getKey());
  assertEquals(1, argument.getValue().getPartNumber());
  assertEquals(40, argument.getValue().getPartSize());
  assertEquals("123", argument.getValue().getUploadId());

  /*
   * Convert the actual InputStream from the client into a ByteArrayOutputStream which can be read
   * and verified.
   */
  byte[] actualBytes = captured.toByteArray();
  byte[] expectedBytes =
      {66, 90, 104, 57, 49, 65, 89, 38, 83, 89, 118, -10, -77, -27, 0, 0, 0, -63, 0, 0, 16, 1, 0,
          -96, 0, 48, -52, 12, -62, 12, 46, -28, -118, 112, -95, 32, -19, -19, 103, -54};

  assertArrayEquals(expectedBytes, actualBytes);
}
 

@Test
public void messageWorkerUsesSameFailureActionForStubAndThrownException() throws Exception {
    // route multicast message to two recipients
    // call failureAction in first stub call
    // throw exception in second stub call
    // make sure that the message is rescheduled only once
    // (multiple executions of the same failure action reschedule only in the first call, further calls are just logged)
    ChannelAddress receiverAddress1 = new ChannelAddress("http://testUrl", "channelId1");
    ChannelAddress receiverAddress2 = new ChannelAddress("http://testUrl", "channelId2");
    prepareMulticastForMultipleAddresses(receiverAddress1, receiverAddress2);
    ImmutableMessage immutableMessage = joynrMessage.getImmutableMessage();

    IMessagingStub messagingStubMock1 = mock(IMessagingStub.class);
    IMessagingStub messagingStubMock2 = mock(IMessagingStub.class);
    reset(middlewareMessagingStubFactoryMock);
    doReturn(messagingStubMock1).when(middlewareMessagingStubFactoryMock).create(receiverAddress1);
    doReturn(messagingStubMock2).when(middlewareMessagingStubFactoryMock).create(receiverAddress2);

    Answer<Void> stubAnswer = new Answer<Void>() {
        private volatile int callCount = 0;

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            callCount++;
            FailureAction failureAction = invocation.getArgumentAt(2, FailureAction.class);
            switch (callCount) {
            case 1:
                failureAction.execute(new JoynrDelayMessageException(32, "first stub call, failureAction"));
                break;
            case 2:
                throw new JoynrDelayMessageException(32, "first stub call, thrown exception");
            case 3:
                // second stub call of stub 1
                break;
            case 4:
                // second stub call of stub 2
                break;
            default:
                fail("expected no more calls");
            }
            return null;
        }
    };
    doAnswer(stubAnswer).when(messagingStubMock1)
                        .transmit(eq(immutableMessage), any(SuccessAction.class), any(FailureAction.class));
    doAnswer(stubAnswer).when(messagingStubMock2)
                        .transmit(eq(immutableMessage), any(SuccessAction.class), any(FailureAction.class));

    messageRouter.route(immutableMessage);

    Thread.sleep(1000);

    verify(messagingStubMock1, times(2)).transmit(eq(immutableMessage),
                                                  any(SuccessAction.class),
                                                  any(FailureAction.class));
    verify(messagingStubMock2, times(2)).transmit(eq(immutableMessage),
                                                  any(SuccessAction.class),
                                                  any(FailureAction.class));
    verify(middlewareMessagingStubFactoryMock, times(2)).create(receiverAddress1);
    verify(middlewareMessagingStubFactoryMock, times(2)).create(receiverAddress2);
}
 

@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
    byte[] buffer = invocation.getArgumentAt(1, byte[].class);
    Arrays.fill(buffer, mValue);
    return null;
}
 

/**
 * Test method to activate a grammar with a field scope and a form scope.
 * 
 * @throws Exception
 *             Test failed.
 * @throws JVoiceXMLEvent
 *             Test failed.
 */
@Test(timeout = 5000)
public void testActivateFormFieldGrammars() throws Exception,
        JVoiceXMLEvent {
    final VoiceXmlDocument doc = new VoiceXmlDocument();
    final Vxml vxml = doc.getVxml();
    final Form form = vxml.appendChild(Form.class);
    final Grammar grammar = form.appendChild(Grammar.class);
    grammar.setVersion("1.0");
    grammar.setType(GrammarType.SRGS_XML);
    final Rule rule = grammar.appendChild(Rule.class);
    final OneOf oneof = rule.appendChild(OneOf.class);
    final Item item1 = oneof.appendChild(Item.class);
    item1.addText("visa");
    final Item item2 = oneof.appendChild(Item.class);
    item2.addText("mastercard");
    final Item item3 = oneof.appendChild(Item.class);
    item3.addText("american express");
    final Field field = form.appendChild(Field.class);
    final Grammar fieldGrammar = field.appendChild(Grammar.class);
    fieldGrammar.setVersion("1.0");
    fieldGrammar.setType(GrammarType.SRGS_XML);
    final Rule fieldRule = fieldGrammar.appendChild(Rule.class);
    final OneOf fieldOneof = fieldRule.appendChild(OneOf.class);
    final Item fieldItem1 = fieldOneof.appendChild(Item.class);
    fieldItem1.addText("euro");
    final Item fieldItem2 = fieldOneof.appendChild(Item.class);
    fieldItem2.addText("dollar");
    final Dialog executableForm = new ExecutablePlainForm();
    executableForm.setNode(form);
    final FormInterpretationAlgorithm fia = new FormInterpretationAlgorithm(
            context, interpreter, executableForm);
    final Object lock = new Object();
    final Thread thread = new Thread() {
        @Override
        public void run() {
            try {
                fia.initialize(profile, null);
                fia.mainLoop();
            } catch (JVoiceXMLEvent e) {
                Assert.fail(e.getMessage());
            }
        };
    };
    thread.start();
    final UserInput input = platform.getUserInput();
    final Answer<Integer> answer = new Answer<Integer>() {
        @Override
        public Integer answer(final InvocationOnMock invocation)
                throws Throwable {
            final Collection grammars = invocation.getArgumentAt(0,
                    Collection.class);
            synchronized (lock) {
                lock.notifyAll();
            }
            return grammars.size();
        }
    };
    Mockito.when(
            input.activateGrammars(Mockito
                    .anyCollectionOf(GrammarDocument.class))).then(answer);
    synchronized (lock) {
        lock.wait();
    }
    Mockito.verify(input).activateGrammars(
            Mockito.anyCollectionOf(GrammarDocument.class));
}