diff --git a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java index 3c39aa33..daf89c28 100644 --- a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java +++ b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java @@ -426,6 +426,13 @@ public class Worker implements Runnable { return applicationName; } + /** + * @return the leaseCoordinator + */ + KinesisClientLibLeaseCoordinator getLeaseCoordinator(){ + return leaseCoordinator; + } + /** * Start consuming data from the stream, and pass it to the application record processors. */ diff --git a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java b/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java index 6cc7ef08..037a54b2 100644 --- a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java +++ b/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java @@ -1500,6 +1500,33 @@ public class WorkerTest { Assert.assertTrue(worker.getStreamConfig().getStreamProxy() instanceof KinesisLocalFileProxy); } + @Test + public void testBuilderWithDefaultLeaseManager() { + IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class); + + Worker worker = new Worker.Builder() + .recordProcessorFactory(recordProcessorFactory) + .config(config) + .build(); + + Assert.assertNotNull(worker.getLeaseCoordinator().getLeaseManager()); + } + + @SuppressWarnings("unchecked") + @Test + public void testBuilderWhenLeaseManagerIsSet() { + IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class); + // Create an instance of ILeaseManager for injection and validation + ILeaseManager leaseManager = (ILeaseManager) mock(ILeaseManager.class); + Worker worker = new Worker.Builder() + .recordProcessorFactory(recordProcessorFactory) + .config(config) + .leaseManager(leaseManager) + .build(); + + Assert.assertSame(leaseManager, worker.getLeaseCoordinator().getLeaseManager()); + } + private abstract class InjectableWorker extends Worker { InjectableWorker(String applicationName, IRecordProcessorFactory recordProcessorFactory, KinesisClientLibConfiguration config, StreamConfig streamConfig,