software.amazon.awssdk.core.async.SdkPublisher Java Examples

@Test
public void queryExclusiveStartKey() {
    insertRecords();
    Map<String, AttributeValue> expectedLastEvaluatedKey = new HashMap<>();
    expectedLastEvaluatedKey.put("id", stringValue(KEYS_ONLY_RECORDS.get(7).getId()));
    expectedLastEvaluatedKey.put("sort", numberValue(KEYS_ONLY_RECORDS.get(7).getSort()));
    expectedLastEvaluatedKey.put("gsi_id", stringValue(KEYS_ONLY_RECORDS.get(7).getGsiId()));
    expectedLastEvaluatedKey.put("gsi_sort", numberValue(KEYS_ONLY_RECORDS.get(7).getGsiSort()));

    SdkPublisher<Page<Record>> publisher =
        keysOnlyMappedIndex.query(QueryEnhancedRequest.builder()
                                                      .queryConditional(keyEqualTo(k -> k.partitionValue("gsi-id-value")))
                                                      .exclusiveStartKey(expectedLastEvaluatedKey)
                                                      .build());

    List<Page<Record>> results = drainPublisher(publisher, 1);
    Page<Record> page = results.get(0);
    assertThat(page.items(), is(KEYS_ONLY_RECORDS.subList(8, 10)));
    assertThat(page.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void scanLimit() {
    insertRecords();

    SdkPublisher<Page<Record>> publisher = mappedTable.scan(r -> r.limit(5));
    publisher.subscribe(page -> page.items().forEach(item -> System.out.println(item)));

    List<Page<Record>> results = drainPublisher(publisher, 3);

    Page<Record> page1 = results.get(0);
    Page<Record> page2 = results.get(1);
    Page<Record> page3 = results.get(2);

    assertThat(page1.items(), is(RECORDS.subList(0, 5)));
    assertThat(page1.lastEvaluatedKey(), is(getKeyMap(4)));
    assertThat(page2.items(), is(RECORDS.subList(5, 10)));
    assertThat(page2.lastEvaluatedKey(), is(getKeyMap(9)));
    assertThat(page3.items(), is(empty()));
    assertThat(page3.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void scanAllRecordsWithFilter() {
    insertRecords();
    Map<String, AttributeValue> expressionValues = new HashMap<>();
    expressionValues.put(":min_value", numberValue(3));
    expressionValues.put(":max_value", numberValue(5));
    Expression expression = Expression.builder()
                                      .expression("sort >= :min_value AND sort <= :max_value")
                                      .expressionValues(expressionValues)
                                      .build();

    SdkPublisher<Page<Record>> publisher =
        mappedTable.scan(ScanEnhancedRequest.builder().filterExpression(expression).build());

    List<Page<Record>> results = drainPublisher(publisher, 1);
    Page<Record> page = results.get(0);

    assertThat(page.items(),
               is(RECORDS.stream().filter(r -> r.sort >= 3 && r.sort <= 5).collect(Collectors.toList())));
    assertThat(page.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void prepareCallsEqualToExecuteAttempts() {
    mockSuccessfulResponse();
    AtomicLong prepareCalls = new AtomicLong(0);
    executeAndWaitError(new TestTransformer<SdkResponse, Void>() {
        @Override
        public CompletableFuture<Void> prepare() {
            prepareCalls.incrementAndGet();
            return super.prepare();
        }

        @Override
        public void onStream(SdkPublisher<ByteBuffer> stream) {
            stream.subscribe(new DrainingSubscriber<ByteBuffer>() {
                @Override
                public void onComplete() {
                    transformFuture().completeExceptionally(RetryableException.builder().message("retry me please: " + prepareCalls.get()).build());
                }
            });
        }
    });
    assertThat(prepareCalls.get()).isEqualTo(1 + RETRY_POLICY.numRetries());
}
 

@Test
public void scanLimit() {
    insertRecords();

    SdkPublisher<Page<Record>> publisher = keysOnlyMappedIndex.scan(r -> r.limit(5));

    List<Page<Record>> results = drainPublisher(publisher, 3);

    Page<Record> page1 = results.get(0);
    Page<Record> page2 = results.get(1);
    Page<Record> page3 = results.get(2);

    assertThat(page1.items(), is(KEYS_ONLY_RECORDS.subList(0, 5)));
    assertThat(page1.lastEvaluatedKey(), is(getKeyMap(4)));
    assertThat(page2.items(), is(KEYS_ONLY_RECORDS.subList(5, 10)));
    assertThat(page2.lastEvaluatedKey(), is(getKeyMap(9)));
    assertThat(page3.items(), is(empty()));
    assertThat(page3.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void queryExclusiveStartKey() {
    Map<String, AttributeValue> exclusiveStartKey = new HashMap<>();
    exclusiveStartKey.put("id", stringValue("id-value"));
    exclusiveStartKey.put("sort", numberValue(7));
    insertRecords();
    SdkPublisher<Page<Record>> publisher =
        mappedTable.query(QueryEnhancedRequest.builder()
                                              .queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))
                                              .exclusiveStartKey(exclusiveStartKey)
                                              .build());

    List<Page<Record>> results = drainPublisher(publisher, 1);
    Page<Record> page = results.get(0);
    assertThat(page.items(), is(RECORDS.subList(8, 10)));
    assertThat(page.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void testEventStreamingOperation() {
    CompletableFuture<Void> responseFuture = client.eventStreamOperation(r -> {
            },
            subscriber -> {},
            new EventStreamOperationResponseHandler() {
                @Override
                public void responseReceived(EventStreamOperationResponse response) {
                }

                @Override
                public void onEventStream(SdkPublisher<EventStream> publisher) {
                }

                @Override
                public void exceptionOccurred(Throwable throwable) {
                }

                @Override
                public void complete() {
                }
            });
    responseFuture.cancel(true);
    assertThat(executeFuture.isCompletedExceptionally()).isTrue();
    assertThat(executeFuture.isCancelled()).isTrue();
}
 

@Test
public void notFoundRecordReturnsNull() {
    insertRecords();

    BatchGetItemEnhancedRequest batchGetItemEnhancedRequest = requestWithNotFoundRecord();

    SdkPublisher<BatchGetResultPage> publisher = enhancedAsyncClient.batchGetItem(batchGetItemEnhancedRequest);

    List<BatchGetResultPage> results = drainPublisher(publisher, 1);
    assertThat(results.size(), is(1));

    BatchGetResultPage page = results.get(0);
    List<Record1> record1List = page.resultsForTable(mappedTable1);
    assertThat(record1List.size(), is(1));
    assertThat(record1List.get(0).getId(), is(0));

    List<Record2> record2List = page.resultsForTable(mappedTable2);
    assertThat(record2List.size(), is(2));
    assertThat(record2List, containsInAnyOrder(RECORDS_2.get(0), RECORDS_2.get(1)));
}
 

@Test
public void getRecordsFromMultipleTables() {
    insertRecords();

    SdkPublisher<BatchGetResultPage> publisher = batchGetResultPageSdkPublisherForBothTables();

    List<BatchGetResultPage> results = drainPublisher(publisher, 1);
    assertThat(results.size(), is(1));
    BatchGetResultPage page = results.get(0);
    List<Record1> record1List = page.resultsForTable(mappedTable1);
    assertThat(record1List.size(), is(2));
    assertThat(record1List, containsInAnyOrder(RECORDS_1.get(0), RECORDS_1.get(1)));

    List<Record2> record2List = page.resultsForTable(mappedTable2);
    assertThat(record2List.size(), is(2));
    assertThat(record2List, containsInAnyOrder(RECORDS_2.get(0), RECORDS_2.get(1)));
}
 

@Test
public void queryAllRecordsWithFilter_viaItems() {
    insertRecords();
    Map<String, AttributeValue> expressionValues = new HashMap<>();
    expressionValues.put(":min_value", numberValue(3));
    expressionValues.put(":max_value", numberValue(5));
    Expression expression = Expression.builder()
                                      .expression("#value >= :min_value AND #value <= :max_value")
                                      .expressionValues(expressionValues)
                                      .expressionNames(Collections.singletonMap("#value", "value"))
                                      .build();

    SdkPublisher<Record> publisher =
        mappedTable.query(QueryEnhancedRequest.builder()
                                              .queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))
                                              .filterExpression(expression)
                                              .build()).items();

    List<Record> results = drainPublisher(publisher, 3);

    assertThat(results,
               is(RECORDS.stream().filter(r -> r.sort >= 3 && r.sort <= 5).collect(Collectors.toList())));
}
 

@Test
public void queryAllRecordsWithFilter() {
    insertRecords();
    Map<String, AttributeValue> expressionValues = new HashMap<>();
    expressionValues.put(":min_value", numberValue(3));
    expressionValues.put(":max_value", numberValue(5));
    Expression expression = Expression.builder()
                                      .expression("#value >= :min_value AND #value <= :max_value")
                                      .expressionValues(expressionValues)
                                      .expressionNames(Collections.singletonMap("#value", "value"))
                                      .build();

    SdkPublisher<Page<Record>> publisher =
        mappedTable.query(QueryEnhancedRequest.builder()
                                              .queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))
                                              .filterExpression(expression)
                                              .build());

    List<Page<Record>> results = drainPublisher(publisher, 1);
    Page<Record> page = results.get(0);

    assertThat(page.items(),
               is(RECORDS.stream().filter(r -> r.sort >= 3 && r.sort <= 5).collect(Collectors.toList())));
    assertThat(page.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void customResponseHandler_InterceptorRecievesResponsePojo() throws Exception {
    final CompletableFuture<String> cf = new CompletableFuture<>();
    try (S3AsyncClient asyncWithInterceptor = createClientWithInterceptor(new AssertingExecutionInterceptor())) {
        String result = asyncWithInterceptor
                .getObject(getObjectRequest, new AsyncResponseTransformer<GetObjectResponse, String>() {

                    @Override
                    public CompletableFuture<String> prepare() {
                        return cf;
                    }

                    @Override
                    public void onResponse(GetObjectResponse response) {
                        // POJO returned by modifyResponse should be delivered to the AsyncResponseTransformer
                        assertThat(response.metadata()).hasEntrySatisfying("x-amz-assert",
                                                                           s -> assertThat(s).isEqualTo("injected-value"));
                    }

                    @Override
                    public void onStream(SdkPublisher<ByteBuffer> publisher) {
                        publisher.subscribe(new SimpleSubscriber(b -> {
                        }) {
                            @Override
                            public void onComplete() {
                                super.onComplete();
                                cf.complete("result");
                            }
                        });
                    }

                    @Override
                    public void exceptionOccurred(Throwable throwable) {
                        cf.completeExceptionally(throwable);
                    }
                }).join();
        assertThat(result).isEqualTo("result");
    }
}
 

@Test(timeout = 2000)
public void prepareResetsSubscriberRef() throws InterruptedException {
    CountDownLatch latch = new CountDownLatch(2);
    AtomicBoolean exceptionThrown = new AtomicBoolean(false);

    AsyncResponseTransformer<SdkResponse, Void> transformer =
            EventStreamAsyncResponseTransformer.builder()
                    .eventStreamResponseHandler(
                            onEventStream(p -> {
                                try {
                                    p.subscribe(e -> {});
                                } catch (Throwable t) {
                                    exceptionThrown.set(true);
                                } finally {
                                    latch.countDown();
                                }
                            }))
                    .eventResponseHandler((r, e) -> null)
                    .executor(Executors.newFixedThreadPool(2))
                    .future(new CompletableFuture<>())
                    .build();

    Flowable<ByteBuffer> bytePublisher = Flowable.empty();

    CompletableFuture<Void> transformFuture = transformer.prepare();
    transformer.onStream(SdkPublisher.adapt(bytePublisher));
    transformFuture.join();

    transformFuture = transformer.prepare();
    transformer.onStream(SdkPublisher.adapt(bytePublisher));
    transformFuture.join();

    latch.await();
    assertThat(exceptionThrown).isFalse();
}
 

@Override
public void onEventStream(SdkPublisher<SubscribeToShardEventStream> publisher) {
    publisher.filter(SubscribeToShardEvent.class)
             .subscribe(b -> {
                 log.debug(() -> "sequenceNumber " + b.records() + "_" + id);
                 receivedData.addAll(b.records().stream().map(Record::data).collect(Collectors.toList()));
             });
}
 

private MethodSpec getMethodsSpecForSingleResultKey(String resultKey) {
    MemberModel resultKeyModel = memberModelForResponseMember(resultKey);

    // TODO: Support other types besides List or Map
    if (!(resultKeyModel.isList() || resultKeyModel.isMap())) {
        return null;
    }

    TypeName resultKeyType = getTypeForResultKey(resultKey);

    return MethodSpec.methodBuilder(resultKeyModel.getFluentGetterMethodName())
                     .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                     .returns(ParameterizedTypeName.get(ClassName.get(SdkPublisher.class), resultKeyType))
                     .addCode("$T getIterator = ",
                              ParameterizedTypeName.get(ClassName.get(Function.class),
                                                        responseType(),
                                                        ParameterizedTypeName.get(ClassName.get(Iterator.class),
                                                                                  resultKeyType)))
                     .addCode(getIteratorLambdaBlock(resultKey, resultKeyModel))
                     .addCode("\n")
                     .addStatement("return $1T.builder().$2L(new $3L()).iteratorFunction(getIterator).$4L($4L).build()",
                                   PaginatedItemsPublisher.class, NEXT_PAGE_FETCHER_MEMBER, nextPageFetcherClassName(),
                                   LAST_PAGE_FIELD)
                     .addJavadoc(CodeBlock.builder()
                                          .add("Returns a publisher that can be used to get a stream of data. You need to "
                                               + "subscribe to the publisher to request the stream of data. The publisher "
                                               + "has a helper forEach method that takes in a {@link $T} and then applies "
                                               + "that consumer to each response returned by the service.",
                                               TypeName.get(Consumer.class))
                                          .build())
                     .build();
}
 

/**
 * Invokes the EventStreamOperationWithOnlyInput operation asynchronously.
 *
 * @param eventStreamOperationWithOnlyInputRequest
 * @return A Java Future containing the result of the EventStreamOperationWithOnlyInput operation returned by the
 *         service.<br/>
 *         The CompletableFuture returned by this method can be completed exceptionally with the following
 *         exceptions.
 *         <ul>
 *         <li>SdkException Base class for all exceptions that can be thrown by the SDK (both service and client).
 *         Can be used for catch all scenarios.</li>
 *         <li>SdkClientException If any client side error occurs such as an IO related failure, failure to get
 *         credentials, etc.</li>
 *         <li>JsonException Base class for all service exceptions. Unknown exceptions will be thrown as an instance
 *         of this type.</li>
 *         </ul>
 * @sample JsonAsyncClient.EventStreamOperationWithOnlyInput
 * @see <a href="http://docs.aws.amazon.com/goto/WebAPI/json-service-2010-05-08/EventStreamOperationWithOnlyInput"
 *      target="_top">AWS API Documentation</a>
 */
@Override
public CompletableFuture<EventStreamOperationWithOnlyInputResponse> eventStreamOperationWithOnlyInput(
        EventStreamOperationWithOnlyInputRequest eventStreamOperationWithOnlyInputRequest,
        Publisher<InputEventStreamTwo> requestStream) {
    try {
        eventStreamOperationWithOnlyInputRequest = applySignerOverride(eventStreamOperationWithOnlyInputRequest,
                EventStreamAws4Signer.create());
        JsonOperationMetadata operationMetadata = JsonOperationMetadata.builder().hasStreamingSuccessResponse(false)
                .isPayloadJson(true).build();

        HttpResponseHandler<EventStreamOperationWithOnlyInputResponse> responseHandler = protocolFactory
                .createResponseHandler(operationMetadata, EventStreamOperationWithOnlyInputResponse::builder);

        HttpResponseHandler<AwsServiceException> errorResponseHandler = createErrorResponseHandler(protocolFactory,
                operationMetadata);
        EventStreamTaggedUnionJsonMarshaller eventMarshaller = EventStreamTaggedUnionJsonMarshaller.builder()
                .putMarshaller(InputEvent.class, new InputEventMarshaller(protocolFactory))
                .putMarshaller(InputEventTwo.class, new InputEventTwoMarshaller(protocolFactory)).build();
        SdkPublisher<InputEventStreamTwo> eventPublisher = SdkPublisher.adapt(requestStream);
        Publisher<ByteBuffer> adapted = eventPublisher.map(event -> eventMarshaller.marshall(event)).map(
                AwsClientHandlerUtils::encodeEventStreamRequestToByteBuffer);

        CompletableFuture<EventStreamOperationWithOnlyInputResponse> executeFuture = clientHandler
                .execute(new ClientExecutionParams<EventStreamOperationWithOnlyInputRequest, EventStreamOperationWithOnlyInputResponse>()
                        .withOperationName("EventStreamOperationWithOnlyInput")
                        .withMarshaller(new EventStreamOperationWithOnlyInputRequestMarshaller(protocolFactory))
                        .withAsyncRequestBody(software.amazon.awssdk.core.async.AsyncRequestBody.fromPublisher(adapted))
                        .withResponseHandler(responseHandler).withErrorResponseHandler(errorResponseHandler)
                        .withInput(eventStreamOperationWithOnlyInputRequest));
        return executeFuture;
    } catch (Throwable t) {
        return CompletableFutureUtils.failedFuture(t);
    }
}
 

/**
 * Returns a publisher that can be used to get a stream of data. You need to subscribe to the publisher to request
 * the stream of data. The publisher has a helper forEach method that takes in a {@link java.util.function.Consumer}
 * and then applies that consumer to each response returned by the service.
 */
public final SdkPublisher<SimpleStruct> items() {
    Function<PaginatedOperationWithResultKeyResponse, Iterator<SimpleStruct>> getIterator = response -> {
        if (response != null && response.items() != null) {
            return response.items().iterator();
        }
        return Collections.emptyIterator();
    };
    return PaginatedItemsPublisher.builder().nextPageFetcher(new PaginatedOperationWithResultKeyResponseFetcher())
                                  .iteratorFunction(getIterator).isLastPage(isLastPage).build();
}
 

/**
 * Returns a publisher that can be used to get a stream of data. You need to subscribe to the publisher to request
 * the stream of data. The publisher has a helper forEach method that takes in a {@link java.util.function.Consumer}
 * and then applies that consumer to each response returned by the service.
 */
public final SdkPublisher<SimpleStruct> items() {
    Function<SameTokenPaginationApiResponse, Iterator<SimpleStruct>> getIterator = response -> {
        if (response != null && response.items() != null) {
            return response.items().iterator();
        }
        return Collections.emptyIterator();
    };
    return PaginatedItemsPublisher.builder().nextPageFetcher(new SameTokenPaginationApiResponseFetcher())
            .iteratorFunction(getIterator).isLastPage(isLastPage).build();
}
 

@Override
public void onStream(SdkPublisher<ByteBuffer> publisher) {
  publisher.subscribe(new Subscriber<ByteBuffer>() {

    @Override
    public void onSubscribe(Subscription s) {
      s.request(Long.MAX_VALUE);
    }

    @Override
    public void onNext(ByteBuffer b) {
      int len = b.remaining();
      dst.setBytes(curOffset, b);
      curOffset += len;
    }

    @Override
    public void onError(Throwable error) {
      future.completeExceptionally(error);
    }

    @Override
    public void onComplete() {
      future.complete(null);
    }
  });
}
 

/**
 * Creates a SubscribeToShardResponseHandler the classic way by implementing the interface
 */
// snippet-start:[kinesis.java2.stream_example.custom_handler]
private static CompletableFuture<Void> responseHandlerBuilderClassic(KinesisAsyncClient client, SubscribeToShardRequest request) {
    SubscribeToShardResponseHandler responseHandler = new SubscribeToShardResponseHandler() {

        @Override
        public void responseReceived(SubscribeToShardResponse response) {
            System.out.println("Receieved initial response");
        }

        @Override
        public void onEventStream(SdkPublisher<SubscribeToShardEventStream> publisher) {
            publisher
                    // Filter to only SubscribeToShardEvents
                    .filter(SubscribeToShardEvent.class)
                    // Flat map into a publisher of just records
                    .flatMapIterable(SubscribeToShardEvent::records)
                    // Limit to 1000 total records
                    .limit(1000)
                    // Batch records into lists of 25
                    .buffer(25)
                    // Print out each record batch
                    .subscribe(batch -> System.out.println("Record Batch - " + batch));
        }

        @Override
        public void complete() {
            System.out.println("All records stream successfully");
        }

        @Override
        public void exceptionOccurred(Throwable throwable) {
            System.err.println("Error during stream - " + throwable.getMessage());
        }
    };
    return client.subscribeToShard(request, responseHandler);
}
 

@Test
public void getAsyncServiceCall_makesTheRightCallAndReturnsResponse() {
    ScanRequest scanRequest = ScanRequest.builder().build();
    ScanPublisher mockScanPublisher = mock(ScanPublisher.class);
    when(mockDynamoDbAsyncClient.scanPaginator(any(ScanRequest.class))).thenReturn(mockScanPublisher);

    SdkPublisher<ScanResponse> response = scanOperation.asyncServiceCall(mockDynamoDbAsyncClient)
                                                       .apply(scanRequest);

    assertThat(response, is(mockScanPublisher));
    verify(mockDynamoDbAsyncClient).scanPaginator(scanRequest);
}
 

@Test
public void scanEmpty_viaItems() {
    SdkPublisher<Record> publisher = mappedTable.scan().items();
    List<Record> results = drainPublisher(publisher, 0);

    assertThat(results, is(empty()));
}
 

@Test(timeout = 1000L)
public void handlerExecutionCompletesIndependentlyOfRequestExecution_CompleteAfterStream() {
    mockSuccessfulResponse_NonSignalingStream();
    // Since we never signal any elements on the response stream, if the async client handler waited for the stream to
    // finish before completing the future, this would never return.
    assertThat(execute(new TestTransformer<SdkResponse, Publisher<ByteBuffer>>() {
        @Override
        public void onStream(SdkPublisher<ByteBuffer> stream) {
            transformFuture().complete(stream);
        }
    })).isNotNull();
}
 

@Test
public void scanAllRecordsDefaultSettings_viaItems() {
    insertRecords();

    SdkPublisher<Record> publisher = mappedTable.scan(ScanEnhancedRequest.builder().build()).items();
    List<Record> results = drainPublisher(publisher, 10);

    assertThat(results, is(RECORDS));
}
 

@Test
public void queryBetween() {
    insertRecords();
    Key fromKey = Key.builder().partitionValue("id-value").sortValue(3).build();
    Key toKey = Key.builder().partitionValue("id-value").sortValue(5).build();
    SdkPublisher<Page<Record>> publisher = mappedTable.query(r -> r.queryConditional(QueryConditional.sortBetween(fromKey, toKey)));

    List<Page<Record>> results = drainPublisher(publisher, 1);
    Page<Record> page = results.get(0);

    assertThat(page.items(),
               is(RECORDS.stream().filter(r -> r.sort >= 3 && r.sort <= 5).collect(Collectors.toList())));
    assertThat(page.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void scanAllRecordsDefaultSettings() {
    insertRecords();

    SdkPublisher<Page<Record>> publisher = mappedTable.scan(ScanEnhancedRequest.builder().build());
    List<Page<Record>> results = drainPublisher(publisher, 1);
    Page<Record> page = results.get(0);

    assertThat(page.items(), is(RECORDS));
    assertThat(page.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void queryBetween_viaItems() {
    insertRecords();
    Key fromKey = Key.builder().partitionValue("id-value").sortValue(3).build();
    Key toKey = Key.builder().partitionValue("id-value").sortValue(5).build();
    SdkPublisher<Record> publisher = mappedTable.query(r -> r.queryConditional(QueryConditional.sortBetween(fromKey, toKey))).items();

    List<Record> results = drainPublisher(publisher, 3);

    assertThat(results,
               is(RECORDS.stream().filter(r -> r.sort >= 3 && r.sort <= 5).collect(Collectors.toList())));
}
 

@Test
public void queryEmpty_viaItems() {
    SdkPublisher<Record> publisher =
        mappedTable.query(r -> r.queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))).items();

    List<Record> results = drainPublisher(publisher, 0);
    assertThat(results, is(empty()));
}
 

@Test
public void queryEmpty() {
    SdkPublisher<Page<Record>> publisher =
        mappedTable.query(r -> r.queryConditional(keyEqualTo(k -> k.partitionValue("id-value"))));

    List<Page<Record>> results = drainPublisher(publisher, 1);
    Page<Record> page = results.get(0);

    assertThat(page.items(), is(empty()));
    assertThat(page.lastEvaluatedKey(), is(nullValue()));
}
 

@Test
public void queryLimit() {
    insertRecords();
    SdkPublisher<Page<Record>> publisher =
        mappedTable.query(QueryEnhancedRequest.builder()
                                              .queryConditional(keyEqualTo(k -> k.partitionValue("id-value")))
                                              .limit(5)
                                              .build());

    List<Page<Record>> results = drainPublisher(publisher, 3);
    Page<Record> page1 = results.get(0);
    Page<Record> page2 = results.get(1);
    Page<Record> page3 = results.get(2);

    Map<String, AttributeValue> expectedLastEvaluatedKey1 = new HashMap<>();
    expectedLastEvaluatedKey1.put("id", stringValue("id-value"));
    expectedLastEvaluatedKey1.put("sort", numberValue(4));
    Map<String, AttributeValue> expectedLastEvaluatedKey2 = new HashMap<>();
    expectedLastEvaluatedKey2.put("id", stringValue("id-value"));
    expectedLastEvaluatedKey2.put("sort", numberValue(9));
    assertThat(page1.items(), is(RECORDS.subList(0, 5)));
    assertThat(page1.lastEvaluatedKey(), is(expectedLastEvaluatedKey1));
    assertThat(page2.items(), is(RECORDS.subList(5, 10)));
    assertThat(page2.lastEvaluatedKey(), is(expectedLastEvaluatedKey2));
    assertThat(page3.items(), is(empty()));
    assertThat(page3.lastEvaluatedKey(), is(nullValue()));
}