Java Code Examples for java.util.concurrent.CompletableFuture#complete()

The following examples show how to use java.util.concurrent.CompletableFuture#complete() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: DisableFTEnableOnMethodTest.java    From microprofile-fault-tolerance with Apache License 2.0 6 votes vote down vote up
/**
 * Test whether Bulkhead is enabled on {@code waitWithBulkhead()}
 *
 * @throws InterruptedException interrupted
 * @throws ExecutionException task was aborted
 */
@Test
public void testBulkhead() throws ExecutionException, InterruptedException {
    ExecutorService executor = Executors.newFixedThreadPool(10);

    // Start two executions at once
    CompletableFuture<Void> waitingFuture = new CompletableFuture<>();
    Future<?> result1 = executor.submit(() -> disableClient.waitWithBulkhead(waitingFuture));
    Future<?> result2 = executor.submit(() -> disableClient.waitWithBulkhead(waitingFuture));

    try {
        disableClient.waitForBulkheadExecutions(2);

        // Try to start a third execution. This would throw a BulkheadException if Bulkhead is enabled.
        // Bulkhead is enabled on the method, so expect exception
        Assert.assertThrows(BulkheadException.class, () -> disableClient.waitWithBulkhead(CompletableFuture.completedFuture(null)));
    }
    finally {
        // Clean up executor and first two executions
        executor.shutdown();

        waitingFuture.complete(null);
        result1.get();
        result2.get();
    }
}
 
Example 2
Source File: StreamSegmentContainerTests.java    From pravega with Apache License 2.0 6 votes vote down vote up
/**
 * Triggers a number of metadata cleanups by repeatedly appending to a random new segment until a cleanup task is detected.
 *
 * @param expectedSegmentNames The segments that we are expecting to evict.
 */
CompletableFuture<Void> triggerMetadataCleanup(Collection<String> expectedSegmentNames) {
    String tempSegmentName = getSegmentName(Long.hashCode(System.nanoTime()));
    HashSet<String> remainingSegments = new HashSet<>(expectedSegmentNames);
    CompletableFuture<Void> cleanupTask = Futures.futureWithTimeout(TIMEOUT, this.executor);

    // Inject this callback into the MetadataCleaner callback, which was setup for us in createMetadataCleaner().
    this.metadataCleanupFinishedCallback = evictedSegmentNames -> {
        remainingSegments.removeAll(evictedSegmentNames);
        if (remainingSegments.size() == 0) {
            cleanupTask.complete(null);
        }
    };

    CompletableFuture<Void> af = appendRandomly(tempSegmentName, true, () -> !cleanupTask.isDone());
    Futures.exceptionListener(af, cleanupTask::completeExceptionally);
    return cleanupTask;
}
 
Example 3
Source File: TaskStoreFactoryForTests.java    From pravega with Apache License 2.0 6 votes vote down vote up
@Override
public CompletableFuture<Void> lock(final Resource resource,
                                    final TaskData taskData,
                                    final String owner,
                                    final String tag,
                                    final String oldOwner,
                                    final String oldTag) {
    CompletableFuture<Void> future = super.lock(resource, taskData, owner, tag, oldOwner, oldTag);

    CompletableFuture<Void> lf = latch.get();
    if (lf != null && first.getAndSet(false)) {
        log.debug("Waiting on the second thread to request the lock and complete the future");
        lf.join();
    } else if (lf != null) {
        log.debug("I'm the second thread, completing the future");
        lf.complete(null);
        latch.set(null);
    } else {
        log.debug("Latch is null");
    }

    return future;
}
 
Example 4
Source File: GuavateTest.java    From Strata with Apache License 2.0 6 votes vote down vote up
@Test
public void test_combineFuturesAsList_Void() {
  CompletableFuture<Void> future1 = new CompletableFuture<>();
  future1.complete(null);
  CountDownLatch latch = new CountDownLatch(1);
  CompletableFuture<Void> future2 = CompletableFuture.supplyAsync(() -> {
    try {
      latch.await();
    } catch (InterruptedException ex) {
      // ignore
    }
    return null;
  });
  List<CompletableFuture<Void>> input = ImmutableList.of(future1, future2);

  CompletableFuture<List<Void>> test = Guavate.combineFuturesAsList(input);

  assertThat(test.isDone()).isEqualTo(false);
  latch.countDown();
  List<Void> combined = test.join();
  assertThat(test.isDone()).isEqualTo(true);
  assertThat(combined.size()).isEqualTo(2);
  assertThat(combined.get(0)).isNull();
  assertThat(combined.get(1)).isNull();
}
 
Example 5
Source File: GuavateTest.java    From Strata with Apache License 2.0 6 votes vote down vote up
@Test
public void test_combineFuturesAsMap_exception() {
  CompletableFuture<String> future1 = new CompletableFuture<>();
  future1.complete("A");
  CountDownLatch latch = new CountDownLatch(1);
  CompletableFuture<String> future2 = CompletableFuture.supplyAsync(() -> {
    try {
      latch.await();
    } catch (InterruptedException ex) {
      // ignore
    }
    throw new IllegalStateException("Oops");
  });
  Map<String, CompletableFuture<String>> input = ImmutableMap.of("a", future1, "b", future2);

  CompletableFuture<Map<String, String>> test = Guavate.combineFuturesAsMap(input);

  assertThat(test.isDone()).isEqualTo(false);
  latch.countDown();
  assertThatExceptionOfType(CompletionException.class).isThrownBy(() -> test.join());
  assertThat(test.isDone()).isEqualTo(true);
  assertThat(test.isCompletedExceptionally()).isEqualTo(true);
}
 
Example 6
Source File: FutureUtilsTest.java    From Flink-CEPplus with Apache License 2.0 6 votes vote down vote up
@Test
public void testRunAfterwards() throws Exception {
	final CompletableFuture<Void> inputFuture = new CompletableFuture<>();
	final OneShotLatch runnableLatch = new OneShotLatch();

	final CompletableFuture<Void> runFuture = FutureUtils.runAfterwards(
		inputFuture,
		runnableLatch::trigger);

	assertThat(runnableLatch.isTriggered(), is(false));
	assertThat(runFuture.isDone(), is(false));

	inputFuture.complete(null);

	assertThat(runnableLatch.isTriggered(), is(true));
	assertThat(runFuture.isDone(), is(true));

	// check that this future is not exceptionally completed
	runFuture.get();
}
 
Example 7
Source File: RlpxConnectionTest.java    From besu with Apache License 2.0 5 votes vote down vote up
@Test
public void getPeerConnection_activeOutboundConnection() {
  final Peer peer = createPeer();
  final CompletableFuture<PeerConnection> future = new CompletableFuture<>();
  final RlpxConnection conn = RlpxConnection.outboundConnection(peer, future);
  final PeerConnection peerConnection = peerConnection(peer);
  future.complete(peerConnection);

  assertThat(conn.getPeerConnection()).isEqualTo(peerConnection);
}
 
Example 8
Source File: StreamReplier.java    From azeroth with Apache License 2.0 5 votes vote down vote up
@Override
protected void readContent(ByteBuf in, CompletableFuture<Void> promise) {
    try {
        int before = in.readableBytes();
        sink.write(in);
        int after = in.readableBytes();
        readed += before - after;

        if (readed >= length) {
            promise.complete(null);
        }
    } catch (IOException e) {
        throw new FastdfsException("write response to output error.", e);
    }
}
 
Example 9
Source File: PingTests.java    From nats.java with Apache License 2.0 5 votes vote down vote up
@Test
public void testHandlingPing() throws IOException, InterruptedException,ExecutionException {
    CompletableFuture<Boolean> gotPong = new CompletableFuture<>();

    NatsServerProtocolMock.Customizer pingPongCustomizer = (ts, r,w) -> {
        
        System.out.println("*** Mock Server @" + ts.getPort() + " sending PING ...");
        w.write("PING\r\n");
        w.flush();

        String pong = "";
        
        System.out.println("*** Mock Server @" + ts.getPort() + " waiting for PONG ...");
        try {
            pong = r.readLine();
        } catch(Exception e) {
            gotPong.cancel(true);
            return;
        }

        if (pong.startsWith("PONG")) {
            System.out.println("*** Mock Server @" + ts.getPort() + " got PONG ...");
            gotPong.complete(Boolean.TRUE);
        } else {
            System.out.println("*** Mock Server @" + ts.getPort() + " got something else... " + pong);
            gotPong.complete(Boolean.FALSE);
        }
    };

    try (NatsServerProtocolMock ts = new NatsServerProtocolMock(pingPongCustomizer)) {
        Connection  nc = Nats.connect(ts.getURI());
        try {
            assertTrue("Connected Status", Connection.Status.CONNECTED == nc.getStatus());
            assertTrue("Got pong.", gotPong.get().booleanValue());
        } finally {
            nc.close();
            assertTrue("Closed Status", Connection.Status.CLOSED == nc.getStatus());
        }
    }
}
 
Example 10
Source File: JavaSyncFutureBenchmark.java    From future with Apache License 2.0 5 votes vote down vote up
@Benchmark
public String flatMapPromiseN() throws InterruptedException, ExecutionException {
  CompletableFuture<String> p = new CompletableFuture<>();
  CompletableFuture<String> f = p;
  for (int i = 0; i < N.n; i++)
    f = f.thenCompose(flatMapF);
  p.complete(string);
  return f.get();
}
 
Example 11
Source File: StandardProcessScheduler.java    From nifi with Apache License 2.0 5 votes vote down vote up
/**
 * Starts the given {@link Processor} by invoking its
 * {@link ProcessorNode#start(ScheduledExecutorService, long, long, Supplier, SchedulingAgentCallback, boolean)}
 * method.
 *
 * @see StandardProcessorNode#start(ScheduledExecutorService, long, long, Supplier, SchedulingAgentCallback, boolean)
 */
@Override
public synchronized CompletableFuture<Void> startProcessor(final ProcessorNode procNode, final boolean failIfStopping) {
    final LifecycleState lifecycleState = getLifecycleState(requireNonNull(procNode), true);

    final Supplier<ProcessContext> processContextFactory = () -> new StandardProcessContext(procNode, getControllerServiceProvider(),
        this.encryptor, getStateManager(procNode.getIdentifier()), lifecycleState::isTerminated);

    final CompletableFuture<Void> future = new CompletableFuture<>();
    final SchedulingAgentCallback callback = new SchedulingAgentCallback() {
        @Override
        public void trigger() {
            lifecycleState.clearTerminationFlag();
            getSchedulingAgent(procNode).schedule(procNode, lifecycleState);
            future.complete(null);
        }

        @Override
        public Future<?> scheduleTask(final Callable<?> task) {
            lifecycleState.incrementActiveThreadCount(null);
            return componentLifeCycleThreadPool.submit(task);
        }

        @Override
        public void onTaskComplete() {
            lifecycleState.decrementActiveThreadCount(null);
        }
    };

    LOG.info("Starting {}", procNode);
    procNode.start(componentMonitoringThreadPool, administrativeYieldMillis, processorStartTimeoutMillis, processContextFactory, callback, failIfStopping);
    return future;
}
 
Example 12
Source File: ReplicaWriterTest.java    From waltz with Apache License 2.0 5 votes vote down vote up
@Test
public void testTransactionOutOfOrder1() throws Exception {
    ReplicaId replicaId = new ReplicaId(PARTITION_ID, String.format(TestReplicaSessionManager.CONNECT_STRING_TEMPLATE, 0));
    MockReplicaConnectionFactory connectionFactory = new MockReplicaConnectionFactory(NUM_PARTITIONS);
    connectionFactory.setCurrentSession(PARTITION_ID, SESSION_ID);

    try {
        CompletableFuture<ReplicaConnection> future = new CompletableFuture<>();
        future.complete(connectionFactory.get(replicaId.partitionId, SESSION_ID));

        ReplicaWriter writer = new ReplicaWriter(future);
        writer.open(0L);

        writer.append(0, TestUtils.makeStoreAppendRequests(0L, 10L));

        assertEquals(-1L, connectionFactory.getLowWaterMark(PARTITION_ID, SESSION_ID));
        assertEquals(9L, connectionFactory.getMaxTransactionId(PARTITION_ID, SESSION_ID));
        assertEquals(10L, writer.nextTransactionId());


        writer.append(10L, TestUtils.makeStoreAppendRequests(10L, 20L));

        assertEquals(-1L, connectionFactory.getLowWaterMark(PARTITION_ID, SESSION_ID));
        assertEquals(19L, connectionFactory.getMaxTransactionId(PARTITION_ID, SESSION_ID));
        assertEquals(20L, writer.nextTransactionId());

        try {
            // overlapping transaction id
            writer.append(19L, TestUtils.makeStoreAppendRequests(19L, 30L));
            fail();
        } catch (ReplicaWriterException ex) {
            // OK
        }
        assertEquals(-1L, connectionFactory.getLowWaterMark(PARTITION_ID, SESSION_ID));
        assertEquals(19L, connectionFactory.getMaxTransactionId(PARTITION_ID, SESSION_ID));

    } finally {
        connectionFactory.close();
    }
}
 
Example 13
Source File: LocalTable.java    From samza with Apache License 2.0 5 votes vote down vote up
@Override
public CompletableFuture<V> getAsync(K key, Object ... args) {
  CompletableFuture<V> future = new CompletableFuture<>();
  try {
    future.complete(get(key, args));
  } catch (Exception e) {
    future.completeExceptionally(e);
  }
  return future;
}
 
Example 14
Source File: CompletableFutureComplete.java    From LearningOfThinkInJava with Apache License 2.0 5 votes vote down vote up
public static void main(String[] args) throws Exception{
        final CompletableFuture<Integer> future=new CompletableFuture<>();
        new Thread(new AskThread(future)).start();
        //模拟主线程的计算,然后触发re.get()
        Thread.sleep(100);
        System.out.println("main thread:before complete"+Thread.currentThread().getName()+"  status:"+future);
        future.complete(60);
//        System.out.println("future.get():"+future.get());
        System.out.println("main thread:after complete"+Thread.currentThread().getName()+"  status:"+future);
    }
 
Example 15
Source File: SubtransExecutorTest.java    From pentaho-kettle with Apache License 2.0 5 votes vote down vote up
private void safeAcquirePermits( SubtransExecutor subtransExecutor, int permitCount,
                                 CompletableFuture<Boolean> threadRunning ) {
  try {
    for ( int i = 0; i < permitCount; i++ ) {
      threadRunning.complete( true );
      subtransExecutor.acquireBufferPermit();
    }
  } catch ( InterruptedException e ) {
    fail();
  }
}
 
Example 16
Source File: PostgresWireProtocol.java    From crate with Apache License 2.0 4 votes vote down vote up
private CompletableFuture<?> handleSingleQuery(String query, DelayableWriteChannel channel) {

        CompletableFuture<?> result = new CompletableFuture<>();

        if (query.isEmpty() || ";".equals(query)) {
            Messages.sendEmptyQueryResponse(channel);
            result.complete(null);
            return result;
        }

        try {
            session.parse("", query, Collections.emptyList());
            session.bind("", "", Collections.emptyList(), null);
            DescribeResult describeResult = session.describe('P', "");
            List<Symbol> fields = describeResult.getFields();

            Function<Throwable, Exception> wrapError = SQLExceptions.forWireTransmission(
                getAccessControl.apply(session.sessionContext()));
            CompletableFuture<?> execute;
            if (fields == null) {
                RowCountReceiver rowCountReceiver = new RowCountReceiver(query, channel.bypassDelay(), wrapError);
                execute = session.execute("", 0, rowCountReceiver);
            } else {
                Messages.sendRowDescription(channel, fields, null);
                ResultSetReceiver resultSetReceiver = new ResultSetReceiver(
                    query,
                    channel.bypassDelay(),
                    wrapError,
                    Lists2.map(fields, x -> PGTypes.get(x.valueType())),
                    null
                );
                execute = session.execute("", 0, resultSetReceiver);
            }
            if (execute == null) {
                return session.sync();
            } else {
                channel.delayWritesUntil(execute);
                return execute.thenCompose(ignored -> session.sync());
            }
        } catch (Throwable t) {
            Messages.sendErrorResponse(channel, t);
            result.completeExceptionally(t);
            return result;
        }
    }
 
Example 17
Source File: EthProtocolManagerTest.java    From besu with Apache License 2.0 4 votes vote down vote up
@Test
public void respondToGetNodeData() throws Exception {
  final CompletableFuture<Void> done = new CompletableFuture<>();
  final WorldStateArchive worldStateArchive = protocolContext.getWorldStateArchive();

  try (final EthProtocolManager ethManager =
      EthProtocolManagerTestUtil.create(
          blockchain,
          () -> false,
          protocolContext.getWorldStateArchive(),
          transactionPool,
          EthProtocolConfiguration.defaultConfig())) {
    // Setup node data query

    final List<Bytes> expectedResults = new ArrayList<>();
    final List<Hash> requestedHashes = new ArrayList<>();

    final long startBlock = blockchain.getChainHeadBlockNumber() - 5;
    final int blockCount = 2;
    for (int i = 0; i < blockCount; i++) {
      final BlockHeader header = blockchain.getBlockHeader(startBlock + i).get();
      requestedHashes.add(header.getStateRoot());
      expectedResults.add(worldStateArchive.getNodeData(header.getStateRoot()).get());
    }
    final MessageData messageData = GetNodeDataMessage.create(requestedHashes);

    // Define handler to validate response
    final PeerSendHandler onSend =
        (cap, message, conn) -> {
          if (message.getCode() == EthPV62.STATUS) {
            // Ignore status message
            return;
          }
          assertThat(message.getCode()).isEqualTo(EthPV63.NODE_DATA);
          final NodeDataMessage receiptsMessage = NodeDataMessage.readFrom(message);
          final List<Bytes> nodeData = receiptsMessage.nodeData();
          assertThat(nodeData.size()).isEqualTo(blockCount);
          for (int i = 0; i < blockCount; i++) {
            assertThat(expectedResults.get(i)).isEqualTo(nodeData.get(i));
          }
          done.complete(null);
        };

    // Run test
    final PeerConnection peer = setupPeer(ethManager, onSend);
    ethManager.processMessage(EthProtocol.ETH63, new DefaultMessage(peer, messageData));
    done.get();
  }
}
 
Example 18
Source File: StreamMetadataTasksTest.java    From pravega with Apache License 2.0 4 votes vote down vote up
@Test(timeout = 30000)
public void concurrentCreateStreamTest() {
    TaskMetadataStore taskMetadataStore = spy(TaskStoreFactory.createZKStore(zkClient, executor));

    StreamMetadataTasks metadataTask = new StreamMetadataTasks(streamStorePartialMock, bucketStore, taskMetadataStore, 
            SegmentHelperMock.getSegmentHelperMock(), executor, "host", 
            new GrpcAuthHelper(authEnabled, "key", 300), requestTracker);

    final ScalingPolicy policy = ScalingPolicy.fixed(2);

    String stream = "concurrent";
    final StreamConfiguration config = StreamConfiguration.builder().scalingPolicy(policy).build();
    
    CompletableFuture<Void> createStreamCalled = new CompletableFuture<>();
    CompletableFuture<Void> waitOnCreateStream = new CompletableFuture<>();
    
    doAnswer(x -> {
        createStreamCalled.complete(null);
        waitOnCreateStream.join();
        return x.callRealMethod();
    }).when(streamStorePartialMock).createStream(anyString(), anyString(), any(), anyLong(), any(), any());
    
    CompletableFuture<Controller.CreateStreamStatus.Status> createStreamFuture1 = metadataTask.createStreamRetryOnLockFailure(
            SCOPE, stream, config, System.currentTimeMillis(), 10);

    // wait until create stream is called. let create stream be blocked on `wait` future. 
    createStreamCalled.join();

    // start a new create stream with 1 retries. this should throw lock failed exception
    // second request should fail with LockFailedException as we have not asked for a retry. 
    AssertExtensions.assertFutureThrows("Lock Failed Exception should be thrown", 
            metadataTask.createStreamRetryOnLockFailure(SCOPE, stream, config, System.currentTimeMillis(), 1), 
            e -> Exceptions.unwrap(e) instanceof LockFailedException);

    CompletableFuture<Void> signalLockFailed = new CompletableFuture<>();
    CompletableFuture<Void> waitOnLockFailed = new CompletableFuture<>();

    // first time lock failed exception is thrown, we will complete `signalLockFailed` to indicate lock failed exception is 
    // being thrown.
    // For all subsequent times we will wait on waitOnLockFailed future.  
    doAnswer(x -> {
        CompletableFuture<Void> future = (CompletableFuture<Void>) x.callRealMethod();
        return future.exceptionally(e -> {
            if (Exceptions.unwrap(e) instanceof LockFailedException) {
                if (!signalLockFailed.isDone()) {
                    signalLockFailed.complete(null);
                } else {
                    waitOnLockFailed.join();
                }
            }
            throw new CompletionException(e);
        });
    }).when(taskMetadataStore).lock(any(), any(), anyString(), anyString(), any(), any());

    // start a new create stream with retries. 
    CompletableFuture<Controller.CreateStreamStatus.Status> createStreamFuture2 =
            metadataTask.createStreamRetryOnLockFailure(SCOPE, stream, config, System.currentTimeMillis(), 10);

    // wait until lock failed exception is thrown
    signalLockFailed.join();
    
    // now complete first createStream request
    waitOnCreateStream.complete(null);

    assertEquals(createStreamFuture1.join(), Controller.CreateStreamStatus.Status.SUCCESS);
    
    // now let the lock failed exception be thrown for second request for subsequent retries
    waitOnLockFailed.complete(null);

    // second request should also succeed now but with stream exists
    assertEquals(createStreamFuture2.join(), Controller.CreateStreamStatus.Status.STREAM_EXISTS);
}
 
Example 19
Source File: SlotSharingManagerTest.java    From flink with Apache License 2.0 4 votes vote down vote up
/**
 * Tests that the logical task slot futures are completed once the slot context
 * future is completed.
 */
@Test
public void testSlotContextFutureCompletion() throws Exception {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	final SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	final SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway());

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();
	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		slotContextFuture,
		new SlotRequestId());

	Locality locality1 = Locality.LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot1 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality1);

	Locality locality2 = Locality.HOST_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot2 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality2);

	CompletableFuture<LogicalSlot> logicalSlotFuture1 = singleTaskSlot1.getLogicalSlotFuture();
	CompletableFuture<LogicalSlot> logicalSlotFuture2 = singleTaskSlot2.getLogicalSlotFuture();
	assertFalse(logicalSlotFuture1.isDone());
	assertFalse(logicalSlotFuture2.isDone());

	slotContextFuture.complete(slotContext);

	assertTrue(logicalSlotFuture1.isDone());
	assertTrue(logicalSlotFuture2.isDone());

	final LogicalSlot logicalSlot1 = logicalSlotFuture1.get();
	final LogicalSlot logicalSlot2 = logicalSlotFuture2.get();

	assertEquals(logicalSlot1.getAllocationId(), slotContext.getAllocationId());
	assertEquals(logicalSlot2.getAllocationId(), slotContext.getAllocationId());
	assertEquals(locality1, logicalSlot1.getLocality());
	assertEquals(locality2, logicalSlot2.getLocality());

	Locality locality3 = Locality.NON_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot3 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality3);

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = singleTaskSlot3.getLogicalSlotFuture();

	assertTrue(logicalSlotFuture3.isDone());
	LogicalSlot logicalSlot3 = logicalSlotFuture3.get();

	assertEquals(locality3, logicalSlot3.getLocality());
	assertEquals(slotContext.getAllocationId(), logicalSlot3.getAllocationId());
}
 
Example 20
Source File: PulsarKafkaProducerTest.java    From pulsar with Apache License 2.0 4 votes vote down vote up
@Test
public void testPulsarKafkaSendAvro() throws PulsarClientException {
    // Arrange
    PulsarClient mockClient = mock(PulsarClient.class);
    ProducerBuilder mockProducerBuilder = mock(ProducerBuilder.class);
    org.apache.pulsar.client.api.Producer mockProducer = mock(org.apache.pulsar.client.api.Producer.class);
    ClientBuilder mockClientBuilder = mock(ClientBuilder.class);
    CompletableFuture mockPartitionFuture = new CompletableFuture();
    CompletableFuture mockSendAsyncFuture = new CompletableFuture();
    TypedMessageBuilder mockTypedMessageBuilder = mock(TypedMessageBuilderImpl.class);

    mockPartitionFuture.complete(new ArrayList<>());
    mockSendAsyncFuture.complete(new MessageIdImpl(1, 1, 1));
    doReturn(mockClientBuilder).when(mockClientBuilder).serviceUrl(anyString());
    doReturn(mockClientBuilder).when(mockClientBuilder).keepAliveInterval(anyInt(), any(TimeUnit.class));
    doReturn(mockClient).when(mockClientBuilder).build();
    doReturn(mockPartitionFuture).when(mockClient).getPartitionsForTopic(anyString());
    doReturn(mockProducerBuilder).when(mockProducerBuilder).topic(anyString());
    doReturn(mockProducerBuilder).when(mockProducerBuilder).clone();
    doReturn(mockProducer).when(mockProducerBuilder).create();
    doReturn(mockTypedMessageBuilder).when(mockProducer).newMessage();
    doReturn(mockSendAsyncFuture).when(mockTypedMessageBuilder).sendAsync();
    PowerMockito.mockStatic(PulsarClientKafkaConfig.class);
    PowerMockito.mockStatic(PulsarProducerKafkaConfig.class);
    when(PulsarClientKafkaConfig.getClientBuilder(any(Properties.class))).thenReturn(mockClientBuilder);
    when(PulsarProducerKafkaConfig.getProducerBuilder(any(PulsarClient.class), any(Properties.class))).thenReturn(mockProducerBuilder);

    Properties properties = new Properties();
    properties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
    properties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
    properties.put(ProducerConfig.PARTITIONER_CLASS_CONFIG, DefaultPartitioner.class);
    properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, Arrays.asList("pulsar://localhost:6650"));
    properties.put(ProducerConfig.CONNECTIONS_MAX_IDLE_MS_CONFIG, "1000000");
    properties.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, "1000000");

    AvroSchema<Bar> barSchema = AvroSchema.of(SchemaDefinition.<Bar>builder().withPojo(Bar.class).build());
    AvroSchema<Foo> fooSchema = AvroSchema.of(SchemaDefinition.<Foo>builder().withPojo(Foo.class).build());
    // Act
    PulsarKafkaProducer<Foo, Bar> pulsarKafkaProducer = new PulsarKafkaProducer<>(properties, fooSchema, barSchema);

    Bar bar = new Bar();
    bar.setField1(true);

    Foo foo = new Foo();
    foo.setField1("field1");
    foo.setField2("field2");
    foo.setField3(3);

    pulsarKafkaProducer.send(new ProducerRecord<>("topic", 1, foo, bar));

    // Verify
    verify(mockTypedMessageBuilder).sendAsync();
}