Source:
<code><START> public void testPreloadSplitsForRewindableSplitSource() <END> throws Exception { ConnectorSession session = new TestingConnectorSession( new HiveSessionProperties( new HiveClientConfig().setUseRewindableSplitSource(true), new OrcFileWriterConfig(), new ParquetFileWriterConfig()).getSessionProperties()); HiveSplitSource hiveSplitSource = HiveSplitSource.bucketedRewindable( session, "database", "table", TupleDomain.all(), 10, new DataSize(1, MEGABYTE), new TestingHiveSplitLoader(), EXECUTOR, new CounterStat()); for (int i = 0; i < 10; i++) { hiveSplitSource.addToQueue(new TestSplit(i, OptionalInt.of(0))); assertEquals(hiveSplitSource.getBufferedInternalSplitCount(), i + 1); } SettableFuture<List<ConnectorSplit>> splits = SettableFuture.create(); CountDownLatch started = new CountDownLatch(1); Thread getterThread = new Thread(() -> { try { started.countDown(); List<ConnectorSplit> batch = getSplits(hiveSplitSource, OptionalInt.of(0), 10); splits.set(batch); } catch (Throwable e) { splits.setException(e); } }); getterThread.start(); try { assertTrue(started.await(1, SECONDS)); MILLISECONDS.sleep(200); assertFalse(splits.isDone()); hiveSplitSource.noMoreSplits(); List<ConnectorSplit> connectorSplits = splits.get(800, MILLISECONDS); assertEquals(connectorSplits.size(), 10); for (int i = 0; i < 10; i++) { assertEquals(((HiveSplit) connectorSplits.get(i)).getSchema().getProperty("id"), Integer.toString(i)); } } finally { getterThread.interrupt(); } }</code><technical_language>Some feedback this test: > assertTrue(started.await(1, SECONDS)); higher timeouts (e.g.: 10 seconds) avoid flaky behaviour > List<ConnectorSplit> connectorSplits = splits.get(800, MILLISECONDS); Ditto</technical_language>

Target:
public void testPreloadSplitsForRewindableSplitSource() throws Exception { ConnectorSession session = new TestingConnectorSession( new HiveSessionProperties( new HiveClientConfig().setUseRewindableSplitSource(true), new OrcFileWriterConfig(), new ParquetFileWriterConfig()).getSessionProperties()); HiveSplitSource hiveSplitSource = HiveSplitSource.bucketedRewindable( session, "database", "table", TupleDomain.all(), 10, new DataSize(1, MEGABYTE), new TestingHiveSplitLoader(), EXECUTOR, new CounterStat()); for (int i = 0; i < 10; i++) { hiveSplitSource.addToQueue(new TestSplit(i, OptionalInt.of(0))); assertEquals(hiveSplitSource.getBufferedInternalSplitCount(), i + 1); } SettableFuture<List<ConnectorSplit>> splits = SettableFuture.create(); CountDownLatch started = new CountDownLatch(1); Thread getterThread = new Thread(() -> { try { started.countDown(); List<ConnectorSplit> batch = getSplits(hiveSplitSource, OptionalInt.of(0), 10); splits.set(batch); } catch (Throwable e) { splits.setException(e); } }); getterThread.start(); try { assertTrue(started.await(10, SECONDS)); MILLISECONDS.sleep(200); assertFalse(splits.isDone()); hiveSplitSource.noMoreSplits(); List<ConnectorSplit> connectorSplits = splits.get(10, SECONDS); assertEquals(connectorSplits.size(), 0); assertFalse(hiveSplitSource.isFinished()); connectorSplits = getSplits(hiveSplitSource, OptionalInt.of(0), 10); for (int i = 0; i < 10; i++) { assertEquals(((HiveSplit) connectorSplits.get(i)).getSchema().getProperty("id"), Integer.toString(i)); } assertTrue(hiveSplitSource.isFinished()); } finally { getterThread.interrupt(); } }