Fix ManagedLedgerTest.testManagedLedgerWithReadEntryTimeOut

Author: lhotari

managed-ledger/src/test/java/org/apache/bookkeeper/mledger/impl/ManagedLedgerTest.java

@@ -94,6 +94,7 @@
 import org.apache.bookkeeper.client.LedgerHandle;
 import org.apache.bookkeeper.client.PulsarMockBookKeeper;
 import org.apache.bookkeeper.client.PulsarMockLedgerHandle;
+import org.apache.bookkeeper.client.PulsarMockReadHandleInterceptor;
 import org.apache.bookkeeper.client.api.LedgerEntries;
 import org.apache.bookkeeper.client.api.LedgerMetadata;
 import org.apache.bookkeeper.client.api.ReadHandle;
@@ -3133,17 +3134,26 @@ public void testManagedLedgerWithReadEntryTimeOut() throws Exception {
         ManagedLedgerConfig config = new ManagedLedgerConfig().setReadEntryTimeoutSeconds(1);
         ManagedLedgerImpl ledger = (ManagedLedgerImpl) factory.open("timeout_ledger_test", config);
 
-        BookKeeper bk = mock(BookKeeper.class);
-        doNothing().when(bk).asyncCreateLedger(anyInt(), anyInt(), anyInt(), any(), any(), any(), any(), any());
+        Position position = ledger.addEntry("entry-1".getBytes());
+
+        // ensure that the read isn't cached
+        factory.getEntryCacheManager().clear();
+
+        bkc.setReadHandleInterceptor(new PulsarMockReadHandleInterceptor() {
+            @Override
+            public CompletableFuture<LedgerEntries> interceptReadAsync(long ledgerId, long firstEntry, long lastEntry,
+                                                                       LedgerEntries entries) {
+                return CompletableFuture.supplyAsync(() -> {
+                    return entries;
+                }, CompletableFuture.delayedExecutor(3, TimeUnit.SECONDS));
+            }
+        });
+
         AtomicReference<ManagedLedgerException> responseException1 = new AtomicReference<>();
         String ctxStr = "timeoutCtx";
-        CompletableFuture<LedgerEntries> entriesFuture = new CompletableFuture<>();
-        ReadHandle ledgerHandle = mock(ReadHandle.class);
-        doReturn(entriesFuture).when(ledgerHandle).readAsync(PositionFactory.EARLIEST.getLedgerId(),
-                PositionFactory.EARLIEST.getEntryId());
 
         // (1) test read-timeout for: ManagedLedger.asyncReadEntry(..)
-        ledger.asyncReadEntry(ledgerHandle, PositionFactory.EARLIEST, new ReadEntryCallback() {
+        ledger.asyncReadEntry(position, new ReadEntryCallback() {
             @Override
             public void readEntryComplete(Entry entry, Object ctx) {
                 responseException1.set(null);
@@ -3155,18 +3165,20 @@ public void readEntryFailed(ManagedLedgerException exception, Object ctx) {
                 responseException1.set(exception);
             }
         }, ctxStr);
-        ledger.asyncCreateLedger(bk, config, null, (rc, lh, ctx) -> {}, Collections.emptyMap());
-        retryStrategically((test) -> responseException1.get() != null, 5, 1000);
-        assertNotNull(responseException1.get());
-        assertTrue(responseException1.get().getMessage()
-                .startsWith(BKException.getMessage(BKException.Code.TimeoutException)));
 
-        // (2) test read-timeout for: ManagedLedger.asyncReadEntry(..)
-        AtomicReference<ManagedLedgerException> responseException2 = new AtomicReference<>();
-        Position readPositionRef = PositionFactory.EARLIEST;
-        ManagedCursorImpl cursor = new ManagedCursorImpl(bk, ledger, "cursor1");
-        OpReadEntry opReadEntry = OpReadEntry.create(cursor, readPositionRef, 1, new ReadEntriesCallback() {
+        Awaitility.await().untilAsserted(() -> {
+            assertNotNull(responseException1.get());
+            assertTrue(responseException1.get().getMessage()
+                    .startsWith(BKException.getMessage(BKException.Code.TimeoutException)));
+        });
 
+        // ensure that the read isn't cached
+        factory.getEntryCacheManager().clear();
+
+        // (2) test read-timeout for: ManagedCursor.asyncReadEntries(..)
+        AtomicReference<ManagedLedgerException> responseException2 = new AtomicReference<>();
+        ManagedCursor cursor = ledger.openCursor("cursor1", InitialPosition.Earliest);
+        cursor.asyncReadEntries(1, new ReadEntriesCallback() {
             @Override
             public void readEntriesComplete(List<Entry> entries, Object ctx) {
             }
@@ -3176,16 +3188,13 @@ public void readEntriesFailed(ManagedLedgerException exception, Object ctx) {
                 assertEquals(ctxStr, (String) ctx);
                 responseException2.set(exception);
             }
+        }, ctxStr, PositionFactory.LATEST);
 
-        }, null, PositionFactory.LATEST, null);
-        ledger.asyncReadEntry(ledgerHandle, PositionFactory.EARLIEST.getEntryId(), PositionFactory.EARLIEST.getEntryId(),
-                opReadEntry, ctxStr);
-        retryStrategically((test) -> {
-            return responseException2.get() != null;
-        }, 5, 1000);
-        assertNotNull(responseException2.get());
-        assertTrue(responseException2.get().getMessage()
-                .startsWith(BKException.getMessage(BKException.Code.TimeoutException)));
+        Awaitility.await().untilAsserted(() -> {
+            assertNotNull(responseException2.get());
+            assertTrue(responseException2.get().getMessage()
+                    .startsWith(BKException.getMessage(BKException.Code.TimeoutException)));
+        });
 
         ledger.close();
     }