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

Answer<Void> writeAnswer(final boolean closeIfNoFrame) {
    return new Answer<Void>() {
        @Override
        public Void answer(InvocationOnMock in) throws Throwable {
            Http2Stream stream = in.getArgument(0);
            int numBytes = in.getArgument(1);
            TestStreamByteDistributorStreamState state = stateMap.get(stream.id());
            state.pendingBytes -= numBytes;
            state.hasFrame = state.pendingBytes > 0;
            state.isWriteAllowed = state.hasFrame;
            if (closeIfNoFrame && !state.hasFrame) {
                stream.close();
            }
            distributor.updateStreamableBytes(state);
            return null;
        }
    };
}
 

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

    final byte[] receivedMessage = invocation.getArgument(1);
    final byte[] expectedBinary = expectedMessage.getBytes();
    final int offset = invocation.getArgument(2);
    final int length = invocation.getArgument(3);
    assertNotNull(receivedMessage);
    assertEquals(expectedBinary.length, receivedMessage.length);
    assertEquals(expectedMessage, new String(receivedMessage));
    assertEquals(0, offset);
    assertEquals(expectedBinary.length, length);
    latch.countDown();
    return null;
}
 

@Override
public InputReader answer(final InvocationOnMock invocationOnMock) throws Throwable {
  final List<CompletableFuture<DataUtil.IteratorWithNumBytes>> inputFutures = new ArrayList<>(SOURCE_PARALLELISM);
  final int elementsPerSource = DATA_SIZE / SOURCE_PARALLELISM;
  for (int i = 0; i < SOURCE_PARALLELISM; i++) {
    inputFutures.add(CompletableFuture.completedFuture(
      DataUtil.IteratorWithNumBytes.of(elements.subList(i * elementsPerSource, (i + 1) * elementsPerSource)
        .iterator())));
  }
  final InputReader inputReader = mock(InputReader.class);
  final IRVertex srcVertex = (IRVertex) invocationOnMock.getArgument(1);
  srcVertex.setProperty(ParallelismProperty.of(SOURCE_PARALLELISM));
  when(inputReader.getSrcIrVertex()).thenReturn(srcVertex);
  when(inputReader.read()).thenReturn(inputFutures);
  when(inputReader.getProperties()).thenReturn(new ExecutionPropertyMap<>(""));
  return inputReader;
}
 

@Test
public final void verifyPutSecretValueStoresKey() throws IOException {
    // given
    final String expected = "expected";
    final Key key = mock(Key.class);
    final Answer<?> answer = new Answer<Void>() {
        public Void answer(final InvocationOnMock invocation) throws Throwable {
            final OutputStream stream = invocation.getArgument(0);
            stream.write(expected.getBytes("UTF-8"));
            return null;
        }
    };
    doAnswer(answer).when(key).writeTo(any(OutputStream.class));

    // when
    manager.putSecretValue("secret", "version", key, PREVIOUS);

    // then
    final PutSecretValueRequest request = new PutSecretValueRequest();
    request.setSecretId("secret");
    request.setClientRequestToken("version");
    request.setVersionStages(singleton("AWSPREVIOUS"));
    request.setSecretBinary(ByteBuffer.wrap(expected.getBytes("UTF-8")));
    verify(delegate).putSecretValue(eq(request));
}
 

@Override
public Object answer(InvocationOnMock invocation)
        throws Throwable {
    final Object arg = invocation.getArgument(0);
    final Runnable runner;
    if (arg instanceof Runnable) {
        runner = (Runnable) arg;
    } else {
        runner = () -> {
            try {
                ((Callable<?>) arg).call();
            } catch (Exception e) {
                fail(e);
            }
        };
    }
    runner.run();
    return null;
}
 

@Override
public Object answer(InvocationOnMock invocationOnMock) {
    File file = invocationOnMock.getArgument(0);
    MediaFile mediaFile = new MediaFile();
    mediaFile.setPath(file.getPath());
    return mediaFile;
}
 

private Answer<Void> writeAnswer() {
    return new Answer<Void>() {
        @Override
        public Void answer(InvocationOnMock in) throws Throwable {
            Http2Stream stream = in.getArgument(0);
            int numBytes = in.getArgument(1);
            TestStreamByteDistributorStreamState state = stateMap.get(stream.id());
            state.pendingBytes -= numBytes;
            state.hasFrame = state.pendingBytes > 0;
            distributor.updateStreamableBytes(state);
            return null;
        }
    };
}
 

private FetchBlockTask createMockTask(final InvocationOnMock invocationOnMock) {
  Bytes32 blockRoot = invocationOnMock.getArgument(1);
  final FetchBlockTask task = mock(FetchBlockTask.class);

  lenient().when(task.getBlockRoot()).thenReturn(blockRoot);
  lenient().when(task.getNumberOfRetries()).thenReturn(0);
  final SafeFuture<FetchBlockResult> future = new SafeFuture<>();
  lenient().when(task.run()).thenReturn(future);
  taskFutures.add(future);

  tasks.add(task);

  return task;
}
 

@Test
public final void verifyPutSecretValueStoresKeys() throws IOException {
    // given
    final String expected = "expected";
    final Key key0 = mock(Key.class);
    final Key key1 = mock(Key.class);
    final Answer<?> answer = new Answer<Void>() {
        public Void answer(final InvocationOnMock invocation) throws Throwable {
            final OutputStream stream = invocation.getArgument(0);
            stream.write(expected.getBytes("UTF-8"));
            return null;
        }
    };
    doAnswer(answer).when(key0).writeTo(any(OutputStream.class));
    doAnswer(answer).when(key1).writeTo(any(OutputStream.class));

    // when
    manager.putSecretValue("secret", "version", asList(key0, key1), PREVIOUS);

    // then
    final PutSecretValueRequest request = new PutSecretValueRequest();
    request.setSecretId("secret");
    request.setClientRequestToken("version");
    request.setVersionStages(singleton("AWSPREVIOUS"));
    request.setSecretBinary(ByteBuffer.wrap((expected + expected).getBytes("UTF-8")));
    verify(delegate).putSecretValue(eq(request));
}
 

@Override
public Stream<JobPosition> answer(InvocationOnMock invocation) throws Throwable {
    Person person = invocation.getArgument(0);
    
    if(person.getName().equals("Peter")) {
        return Stream.<JobPosition>builder().add(new JobPosition("Teacher")).build();
    } 
    
    return Stream.empty();
}
 

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

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

@Override
public MongoDbWriter answer(InvocationOnMock invocation) {
    foundBuffer = invocation.getArgument(1);
    foundFactory = invocation.getArgument(4);
    timesCalled++;
    return writer;
}
 

protected Object assertConnectedEvent(CountDownLatch latch, AtomicReference<String> sessionIdRef, InvocationOnMock invocation) {
    final WebSocketSessionInfo sessionInfo = invocation.getArgument(0);
    assertNotNull(sessionInfo.getLocalAddress());
    assertNotNull(sessionInfo.getRemoteAddress());
    assertNotNull(sessionInfo.getSessionId());
    assertEquals(isSecure(), sessionInfo.isSecure());
    sessionIdRef.set(sessionInfo.getSessionId());
    latch.countDown();
    return null;
}
 

static Threading mockThreading() {
    Threading mock = mock(Threading.class);
    Answer executeAnswer = new Answer() {
        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            Runnable runnable = invocation.getArgument(0);
            runnable.run();
            return null;
        }
    };
    doAnswer(executeAnswer).when(mock).runOnMainThread(any(Runnable.class));
    doAnswer(executeAnswer).when(mock).runInBackground(any(Runnable.class));
    return mock;
}
 

@Override
public Object answer(InvocationOnMock invocation)
        throws Throwable {
    activeCount.incrementAndGet();
    try {
        Runnable runnable = invocation.getArgument(0);
        counter.incrementAndGet();
        runnable.run();
    } finally {
        activeCount.decrementAndGet();
    }
    return null;
}
 

private Cluster saveCluster(InvocationOnMock i) {
    Cluster cluster = i.getArgument(0);
    cluster.setClusterPertain(getClusterPertain());
    cluster.setId(CLUSTER_ID);
    return cluster;
}
 

public Integer answer(InvocationOnMock invocationOnMock) throws Throwable {
    byte[] buffer = invocationOnMock.getArgument(0);
    System.arraycopy(value, 0, buffer, 0, value.length);
    return value.length;
}
 

/**
 * Test the implementation of {@link BatchOptions#consistency(InfluxDB.ConsistencyLevel)} }.
 * @throws InterruptedException
 */
@Test
public void testConsistency() throws InterruptedException {
  String dbName = "write_unittest_" + System.currentTimeMillis();

  InfluxDB spy = spy(influxDB);
  spy.query(new Query("CREATE DATABASE " + dbName));
  spy.setDatabase(dbName);
  try {
    TestAnswer answer = new TestAnswer() {
      @Override
      protected void check(InvocationOnMock invocation) {
        BatchPoints batchPoints = (BatchPoints) invocation.getArgument(0);
        Assertions.assertEquals(params.get("consistencyLevel"), batchPoints.getConsistency());

      }
    };
    doAnswer(answer).when(spy).write(any(BatchPoints.class));

    int n = 0;
    for (final ConsistencyLevel consistencyLevel : ConsistencyLevel.values()) {
      answer.params.put("consistencyLevel", consistencyLevel);
      BatchOptions options = BatchOptions.DEFAULTS.consistency(consistencyLevel).flushDuration(100);
      spy.enableBatch(options);
      Assertions.assertEquals(options.getConsistency(), consistencyLevel);

      writeSomePoints(spy, n, n + 4);
      n += 5;
      Thread.sleep(300);

      verify(spy, atLeastOnce()).write(any(BatchPoints.class));
      QueryResult result = spy.query(new Query("select * from weather", dbName));
      Assertions.assertEquals(n, result.getResults().get(0).getSeries().get(0).getValues().size());


      spy.disableBatch();
    }

  } finally {
    spy.query(new Query("DROP DATABASE " + dbName));
  }
}
 

@Override
public Object answer(InvocationOnMock invocationOnMock) {
    Playlist playlist = invocationOnMock.getArgument(0);
    playlist.setId(id);
    return null;
}
 

private FileSystem setIdForCreatedFileSystemEntry(InvocationOnMock invocation) {
    FileSystem fileSystem = invocation.getArgument(0);
    fileSystem.setId(TEST_FILES_SYSTEM_ID);
    return fileSystem;
}