@Override public Object answer(final InvocationOnMock invocation) throws Throwable { statusCode = invocation.getArgumentAt(0, int.class); return null; } }).when(response).setStatus(Mockito.anyInt());
public Object answer(InvocationOnMock invocation) { StreamObserver<CreateLocalBlockResponse> observer = invocation.getArgumentAt(0, StreamObserver.class); observer.onNext(response); return mStreamObserver; } });
public Object answer(InvocationOnMock invocation) { StreamObserver<CreateLocalBlockResponse> observer = invocation.getArgumentAt(0, StreamObserver.class); observer.onNext(response); return mStreamObserver; } });
public Object answer(InvocationOnMock invocation) { StreamObserver<OpenLocalBlockResponse> observer = invocation.getArgumentAt(0, StreamObserver.class); observer.onNext(response); observer.onCompleted(); return mStreamObserver; } });
@Override public Object answer(final InvocationOnMock invocation) throws Throwable { final String key = invocation.getArgumentAt(0, String.class); final String value = invocation.getArgumentAt(1, String.class); if (value == null) { headersWithNoValue.add(key); } else { headersSent.put(key, value); } return null; } }).when(response).setHeader(Mockito.any(String.class), Mockito.any(String.class));
@Override public Void answer(InvocationOnMock invocation) throws Throwable { System.out.println("Updating task: " + invocation.getArgumentAt(0, Task.class)); updatedTasks.add(invocation.getArgumentAt(0, Task.class)); return null; } }).when(es).updateTask(any(Task.class));
@SuppressWarnings("unchecked") @Override public Void answer(InvocationOnMock invocation) throws Throwable { List<Message> msgs = invocation.getArgumentAt(0, List.class); System.out.println("got messages to publish: " + msgs); messages.addAll(msgs); return null; } }).when(queue).publish(any());
@Override public TestBlockOutStream answer(InvocationOnMock invocation) throws Throwable { Long blockId = invocation.getArgumentAt(0, Long.class); if (!outStreamMap.containsKey(blockId)) { TestBlockOutStream newStream = new TestBlockOutStream(ByteBuffer.allocate(1000), BLOCK_LENGTH); outStreamMap.put(blockId, newStream); } return outStreamMap.get(blockId); } });
/** * Set up gRPC interface mocks. */ @Before public void before() { mClient = mock(BlockWorkerClient.class); mRequestObserver = mock(ClientCallStreamObserver.class); when(mClient.writeBlock(any(StreamObserver.class))).thenAnswer((args) -> { mResponseObserver = args.getArgumentAt(0, ClientResponseObserver.class); return mRequestObserver; }); when(mRequestObserver.isReady()).thenReturn(true); mStream = new GrpcBlockingStream<>(mClient::writeBlock, BUFFER_SIZE, TEST_MESSAGE); }
public Object answer(InvocationOnMock invocation) { StreamObserver<OpenLocalBlockResponse> observer = invocation.getArgumentAt(0, StreamObserver.class); observer.onNext(OpenLocalBlockResponse.newBuilder().setPath("/tmp").build()); observer.onCompleted(); return requestObserver; } });
@Test(expected = ExecutionException.class) public void testCancelRequestAfterTimeout() throws Exception { when(executorService.schedule( any(Runnable.class), eq(WebSocketService.REQUEST_TIMEOUT), eq(TimeUnit.SECONDS))) .then(invocation -> { Runnable runnable = invocation.getArgumentAt(0, Runnable.class); runnable.run(); return null; }); CompletableFuture<Web3ClientVersion> reply = service.sendAsync( request, Web3ClientVersion.class); assertTrue(reply.isDone()); reply.get(); }
@Before public void before() throws Exception { when(webSocketClient.connectBlocking()).thenReturn(true); doAnswer(invocation -> { listener = invocation.getArgumentAt(0, WebSocketListener.class); return null; }).when(webSocketClient).setListener(any()); doAnswer(invocation -> { String message = invocation.getArgumentAt(0, String.class); int requestId = getRequestId(message); sendSubscriptionConfirmation(requestId); return null; }).when(webSocketClient).send(anyString()); webSocketService.connect(); }
@Before public void before() throws Exception { when(webSocketClient.connectBlocking()).thenReturn(true); doAnswer(invocation -> { listener = invocation.getArgumentAt(0, WebSocketListener.class); return null; }).when(webSocketClient).setListener(any()); doAnswer(invocation -> { String message = invocation.getArgumentAt(0, String.class); int requestId = getRequestId(message); sendSubscriptionConfirmation(requestId); return null; }).when(webSocketClient).send(anyString()); webSocketService.connect(); }
@Override public HttpResponse answer( InvocationOnMock invocation ) throws Throwable { HttpClientContext context = invocation.getArgumentAt( 1, HttpClientContext.class ); Credentials cred = context.getCredentialsProvider().getCredentials( new AuthScope( "hname", 1111 ) ); assertEquals( "uname", cred.getUserPrincipal().getName() ); return mockResponse( 200, responseContent ); } } );
SparkSession sparkSession = invocationOnMock.getArgumentAt(0, SparkSession.class); if (sparkSession == null) { return mock(SparkSession.class); when(mockSessionState.getConf()).thenReturn(mockConf); doAnswer(invocationOnMock -> { activeSparkSession.set(invocationOnMock.getArgumentAt(0, SparkSession.class)); return null; }).when(mockSessionState).setSparkSession(any(SparkSession.class));
/** * Checks send after stream is ready. */ @Test public void sendAfterStreamReady() throws Exception { when(mRequestObserver.isReady()).thenReturn(false); doAnswer((args) -> { mOnReadyHandler = args.getArgumentAt(0, Runnable.class); return null; }).when(mRequestObserver).setOnReadyHandler(any(Runnable.class)); mResponseObserver.beforeStart(mRequestObserver); EXECUTOR.submit(() -> { try { // notify ready after a short period of time Thread.sleep(SHORT_TIMEOUT); when(mRequestObserver.isReady()).thenReturn(true); mOnReadyHandler.run(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } }); WriteRequest request = WriteRequest.newBuilder().build(); mStream.send(request, TIMEOUT); verify(mRequestObserver).onNext(request); }
private void testGetInStreamFallback(int workerCount, boolean isPersisted, int[] blockLocations, Map<Integer, Long> failedWorkers, int expectedWorker) throws Exception { WorkerNetAddress[] workers = new WorkerNetAddress[workerCount]; Arrays.setAll(workers, i -> new WorkerNetAddress().setHost(String.format("worker-%d", i))); BlockInfo info = new BlockInfo().setBlockId(BLOCK_ID) .setLocations(Arrays.stream(blockLocations).mapToObj(x -> new BlockLocation().setWorkerAddress(workers[x])).collect(Collectors.toList())); URIStatus dummyStatus = new URIStatus(new FileInfo().setPersisted(isPersisted) .setBlockIds(Collections.singletonList(BLOCK_ID)) .setFileBlockInfos(Collections.singletonList(new FileBlockInfo().setBlockInfo(info)))); BlockLocationPolicy mockPolicy = mock(BlockLocationPolicy.class); when(mockPolicy.getWorker(any())).thenAnswer(arg -> arg .getArgumentAt(0, GetWorkerOptions.class).getBlockWorkerInfos().iterator().next() .getNetAddress()); InStreamOptions options = new InStreamOptions(dummyStatus, OpenFilePOptions.getDefaultInstance(), sConf); options.setUfsReadLocationPolicy(mockPolicy); when(mMasterClient.getBlockInfo(BLOCK_ID)).thenReturn(info); when(mMasterClient.getWorkerInfoList()).thenReturn(Arrays.stream(workers) .map(x -> new WorkerInfo().setAddress(x)).collect((Collectors.toList()))); Map<WorkerNetAddress, Long> failedWorkerAddresses = failedWorkers.entrySet().stream() .map(x -> new AbstractMap.SimpleImmutableEntry<>(workers[x.getKey()], x.getValue())) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); BlockInStream inStream = mBlockStore.getInStream(BLOCK_ID, options, failedWorkerAddresses); assertEquals(workers[expectedWorker], inStream.getAddress()); } }
BlockLocationPolicy mockPolicy = mock(BlockLocationPolicy.class); when(mockPolicy.getWorker(any())).thenAnswer(arg -> arg .getArgumentAt(0, GetWorkerOptions.class).getBlockWorkerInfos().iterator().next() .getNetAddress()); InStreamOptions options =
String queueName = invocation.getArgumentAt(0, String.class); System.out.println(queueName); publishedQueue[0] = queueName; List<Message> messages = invocation.getArgumentAt(1, List.class); publishedMessages.addAll(messages); return null; String messageId = invocation.getArgumentAt(1, String.class); if(publishedMessages.get(0).getId().equals(messageId)) { publishedMessages.remove(0);
String queueName = invocation.getArgumentAt(0, String.class); System.out.println(queueName); queuedTaskCount.incrementAndGet();